#!/usr/bin/env python
# -*- coding: utf-8

import sys
import argparse

import anvio.db as db
import anvio.tables as t
import anvio.dbops as dbops
import anvio.terminal as terminal 

from anvio.errors import ConfigError


run = terminal.Run()
progress = terminal.Progress()

current_version = '16'
next_version    = '17'


def update_profile_db(profile_db_path, just_do_it = False):
    if profile_db_path is None:
        raise ConfigError("No database path is given.")

    # make sure someone is not being funny
    dbops.is_profile_db(profile_db_path)

    # make sure the version is accurate
    profile_db = db.DB(profile_db_path, None, ignore_version = True)
    if str(profile_db.get_version()) != current_version:
        raise ConfigError("Version of this profile database is not %s (hence, this script cannot really do anything)." % current_version)

    if not just_do_it:
        try:
            run.warning("This script will try to upgrade your profile database from v%s to v%s. It happens to be that this\
                         upgrade will only rename every occurrence of `portion_covered` to `detection`. We don't even know\
                         why we called it `portion_covered` in the first place :/ You can just press ENTER to continue.\
                         If you want to cancel the upgrade and think more about it, press CTRL+C now. If you want to avoid\
                         this message the next time, use '--just-do-it'." % (current_version, next_version))
            input("Press ENTER to continue...\n")
        except:
            print()
            sys.exit()

    is_merged = profile_db.get_meta_value('merged')
    tables_in_db = profile_db.get_table_names()
    is_full_profile = 'portion_covered_splits' in tables_in_db or 'atomic_data_splits' in tables_in_db

    run.info('Profile db type', 'Merged' if is_merged else 'Single')
    run.info('Full profile', is_full_profile)

    progress.new("Trying to upgrade the profile database")
    progress.update('...')

    if is_full_profile and is_merged:
        profile_db._exec('ALTER TABLE portion_covered_splits RENAME TO detection_splits;')
        profile_db._exec('ALTER TABLE portion_covered_contigs RENAME TO detection_contigs;')
        profile_db._exec('ALTER TABLE mean_coverage_Q1Q3_splits RENAME TO mean_coverage_Q2Q3_splits;')
        profile_db._exec('ALTER TABLE mean_coverage_Q1Q3_contigs RENAME TO mean_coverage_Q2Q3_contigs;')

        profile_db._exec('DELETE FROM %s WHERE view_id = "portion_covered"' % (t.views_table_name))
        profile_db._exec('INSERT INTO %s VALUES ("detection", "detection_splits")' % t.views_table_name)
        profile_db._exec('DELETE FROM %s WHERE view_id = "mean_coverage_Q1Q3"' % (t.views_table_name))
        profile_db._exec('INSERT INTO %s VALUES ("mean_coverage_Q2Q3", "mean_coverage_Q2Q3_splits")' % t.views_table_name)

    elif is_full_profile and not is_merged:
        profile_db.cursor.execute('ALTER TABLE atomic_data_contigs RENAME TO atomic_data_contigs_TEMP;')
        profile_db.cursor.execute('CREATE TABLE atomic_data_contigs (contig text, std_coverage numeric, mean_coverage numeric, mean_coverage_Q2Q3 numeric, max_normalized_ratio numeric, relative_abundance numeric, detection numeric, abundance numeric, variability numeric, __parent__ text);')
        profile_db.cursor.execute('INSERT INTO atomic_data_contigs(contig, std_coverage, mean_coverage, mean_coverage_Q2Q3, max_normalized_ratio, relative_abundance, detection, abundance, variability, __parent__) SELECT contig, std_coverage, mean_coverage, mean_coverage_Q1Q3, max_normalized_ratio, relative_abundance, portion_covered, abundance, variability, __parent__ FROM atomic_data_contigs_TEMP;')
        profile_db.cursor.execute('DROP TABLE atomic_data_contigs_TEMP;')

        profile_db.cursor.execute('ALTER TABLE atomic_data_splits RENAME TO atomic_data_splits_TEMP;')
        profile_db.cursor.execute('CREATE TABLE atomic_data_splits (contig text, std_coverage numeric, mean_coverage numeric, mean_coverage_Q2Q3 numeric, max_normalized_ratio numeric, relative_abundance numeric, detection numeric, abundance numeric, variability numeric, __parent__ text);')
        profile_db.cursor.execute('INSERT INTO atomic_data_splits(contig, std_coverage, mean_coverage, mean_coverage_Q2Q3, max_normalized_ratio, relative_abundance, detection, abundance, variability, __parent__) SELECT contig, std_coverage, mean_coverage, mean_coverage_Q1Q3, max_normalized_ratio, relative_abundance, portion_covered, abundance, variability, __parent__ FROM atomic_data_splits_TEMP;')
        profile_db.cursor.execute('DROP TABLE atomic_data_splits_TEMP;')

    # update states
    states = profile_db.get_table_as_dict(t.states_table_name)
    for state in states:
        profile_db._exec('DELETE FROM %s WHERE name = "%s"' % (t.states_table_name, state))
        profile_db._exec('INSERT INTO %s VALUES (?,?,?)' % (t.states_table_name), (state, states[state]['content'].replace('portion_covered', 'detection').replace('mean_coverage_Q1Q3', 'mean_coverage_Q2Q3'), states[state]['last_modified']))

    # set the version
    profile_db.remove_meta_key_value_pair('version')
    profile_db.set_version(next_version)

    # bye
    profile_db.disconnect()
    progress.end()


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='A simple script to upgrade profile database from version %s to version %s' % (current_version, next_version))
    parser.add_argument('profile_db', metavar = 'PROFILE_DB', help = "An anvi'o profile database of version %s" % current_version)
    parser.add_argument('--just-do-it', default=False, action="store_true", help = "Do not bother me with warnings")
    args = parser.parse_args()

    try:
        update_profile_db(args.profile_db, just_do_it = args.just_do_it)
    except ConfigError as e:
        print(e)
        sys.exit(-1)
