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

import sys
import argparse

import anvio.db as db
import anvio.dbops as dbops
import anvio.terminal as terminal 

from anvio.errors import ConfigError


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


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

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

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

    is_merged = profile_db.get_meta_value('merged')

    if not just_do_it:
        try:
            run.warning('This script will try to upgrade your profile database. If things go south, you will end up having to\
                         re-profile your data :/ It may be a good idea to back it up first. If you already backed your stuff\
                         or you are certain that it will work, or if you are a lucky person in general, press ENTER to continue.\
                         If you want to cancel the upgrade, press CTRL+C now. If you want to avoid this message the next time,\
                         use "--just-do-it" flag.')
            input("Press ENTER to continue...\n")
        except:
            print()
            sys.exit()

    progress.new("Trying to upgrade the %s profile database" % 'merged' if is_merged else 'single')

    # serious stuff ###################################################################################################
    progress.update('working on the variable nts table ...')
    profile_db.cursor.execute('ALTER TABLE variable_nucleotide_positions RENAME TO variable_nucleotide_positions_TEMP;')
    profile_db.cursor.execute('CREATE TABLE variable_nucleotide_positions (entry_id numeric, sample_id text, split_name text, pos numeric, pos_in_contig numeric, corresponding_gene_call numeric, in_partial_gene_call numeric, in_complete_gene_call numeric, base_pos_in_codon numeric, codon_order_in_gene numeric, coverage numeric, cov_outlier_in_split bool, cov_outlier_in_contig bool, departure_from_reference numeric, competing_nts text, reference text, A numeric, T numeric, C numeric, G numeric, N numeric);')
    profile_db.cursor.execute('INSERT INTO variable_nucleotide_positions(entry_id, sample_id, split_name, pos, pos_in_contig, corresponding_gene_call, in_partial_gene_call, in_complete_gene_call, base_pos_in_codon, codon_order_in_gene, coverage, cov_outlier_in_split, cov_outlier_in_contig, departure_from_reference, competing_nts, reference, A, T, C, G, N) SELECT entry_id, sample_id, split_name, pos, pos_in_contig, corresponding_gene_call, in_partial_gene_call, in_complete_gene_call, base_pos_in_codon, codon_order_in_gene, coverage, cov_outlier_in_split, cov_outlier_in_contig, departure_from_consensus, competing_nts, consensus, A, T, C, G, N FROM variable_nucleotide_positions_TEMP;')
    profile_db.cursor.execute('DROP TABLE variable_nucleotide_positions_TEMP;')

    progress.update('working on the aa freqs table ...')
    profile_db.cursor.execute('ALTER TABLE variable_amino_acid_frequencies RENAME TO variable_amino_acid_frequencies_TEMP;')
    profile_db.cursor.execute('CREATE TABLE variable_amino_acid_frequencies (entry_id numeric, sample_id text, corresponding_gene_call numeric, codon_order_in_gene numeric, reference text, departure_from_reference numeric, coverage numeric, Ala numeric, Arg numeric, Asn numeric, Asp numeric, Cys numeric, Gln numeric, Glu numeric, Gly numeric, His numeric, Ile numeric, Leu numeric, Lys numeric, Met numeric, Phe numeric, Pro numeric, STP numeric, Ser numeric, Thr numeric, Trp numeric, Tyr numeric, Val numeric);')
    profile_db.cursor.execute('INSERT INTO variable_amino_acid_frequencies(entry_id, sample_id, corresponding_gene_call, codon_order_in_gene, reference, departure_from_reference, coverage, Ala, Arg, Asn, Asp, Cys, Gln, Glu, Gly, His, Ile, Leu, Lys, Met, Phe, Pro, STP, Ser, Thr, Trp, Tyr, Val) SELECT entry_id, sample_id, corresponding_gene_call, codon_order_in_gene, consensus, departure_from_consensus, coverage, Ala, Arg, Asn, Asp, Cys, Gln, Glu, Gly, His, Ile, Leu, Lys, Met, Phe, Pro, STP, Ser, Thr, Trp, Tyr, Val FROM variable_amino_acid_frequencies_TEMP;')
    profile_db.cursor.execute('DROP TABLE variable_amino_acid_frequencies_TEMP;')
    # /serious stuff ##################################################################################################

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

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

    run.info_single("Database successfully upgraded to version 14!")


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='A simple script to upgrade profile database to from version 13 version 14')

    parser.add_argument('profile_db', metavar = 'PROFILE_DB', help = 'Profile database (of version 13)')
    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_from_v13_to_v14(args.profile_db, just_do_it = args.just_do_it)
    except ConfigError as e:
        print(e)
        sys.exit(-1)
