import argparse
import os
import json
import logging
import string
import collections

import pandas as pd
import numpy as np
from sqlalchemy import create_engine
import networkx as nx
import matplotlib.pyplot as plt
import seaborn as sns

if os.getcwd().endswith('notebook'):
    os.chdir('..')

from rna_learn.alphabet import CODON_REDUNDANCY
from rna_learn.codon_bias.trna_ai import compute_trna_ai_weights
from rna_learn.codon_bias.graph import load_codon_bias
from rna_learn.codon_bias.distance import compute_codon_bias_distance

sns.set(palette='colorblind', font_scale=1.3)
palette = sns.color_palette()
logging.basicConfig(level=logging.INFO, format="%(asctime)s (%(levelname)s) %(message)s")
logger = logging.getLogger(__name__)

db_path = os.path.join(os.getcwd(), 'data/db/seq.db')
engine = create_engine(f'sqlite+pysqlite:///{db_path}')

codon_to_aa_table = { 
    'ATA':'I', 'ATC':'I', 'ATT':'I', 'ATG':'M', 
    'ACA':'T', 'ACC':'T', 'ACG':'T', 'ACT':'T', 
    'AAC':'N', 'AAT':'N', 'AAA':'K', 'AAG':'K', 
    'AGC':'S', 'AGT':'S', 'AGA':'R', 'AGG':'R',                  
    'CTA':'L', 'CTC':'L', 'CTG':'L', 'CTT':'L', 
    'CCA':'P', 'CCC':'P', 'CCG':'P', 'CCT':'P', 
    'CAC':'H', 'CAT':'H', 'CAA':'Q', 'CAG':'Q', 
    'CGA':'R', 'CGC':'R', 'CGG':'R', 'CGT':'R', 
    'GTA':'V', 'GTC':'V', 'GTG':'V', 'GTT':'V', 
    'GCA':'A', 'GCC':'A', 'GCG':'A', 'GCT':'A', 
    'GAC':'D', 'GAT':'D', 'GAA':'E', 'GAG':'E', 
    'GGA':'G', 'GGC':'G', 'GGG':'G', 'GGT':'G', 
    'TCA':'S', 'TCC':'S', 'TCG':'S', 'TCT':'S', 
    'TTC':'F', 'TTT':'F', 'TTA':'L', 'TTG':'L', 
    'TAC':'Y', 'TAT':'Y', 'TAA':'*', 'TAG':'*', 
    'TGC':'C', 'TGT':'C', 'TGA':'*', 'TGG':'W', 
}

species_query = """
select species_taxid, species, phylum, superkingdom, growth_tmp from species_traits
"""
species_df = pd.read_sql(species_query, engine)
species_taxids = species_df['species_taxid'].values

def compute_mae(y_true, y_pred):
    return np.mean(np.abs(y_true - y_pred))


def predict_numerical_field_from_neighbours(graph, species_taxid, trait, control_for_phylum=True):
    current_node = graph.nodes[species_taxid]
    
    if pd.isnull(current_node[trait]):
        return None, None, None
    
    current_phylum = current_node['phylum']
    values = []
    for key in graph.neighbors(species_taxid):
        node = graph.nodes[key]
        
        if pd.isnull(node[trait]):
            continue
        
        if control_for_phylum and node['phylum'] != current_phylum:
            v = node[trait]
            values.append(v)
        elif not control_for_phylum:
            v = node[trait]
            values.append(v)
            
    if len(values) == 0:
        prediction = None
    else:
        prediction = np.mean(values)

    return current_node[trait], prediction, values

def compute_mae_from_neighbour_predictions(graph, species_taxids, trait, control_for_phylum=True):
    actuals = []
    predictions = []
    for i, species_taxid in enumerate(species_taxids):
        actual, prediction, _ = predict_numerical_field_from_neighbours(
            graph, 
            species_taxid, 
            trait, 
            control_for_phylum,
        )
        
        if actual is None or prediction is None:
            continue
        
        actuals.append(actual)
        predictions.append(prediction)
        
    return compute_mae(np.array(actuals), np.array(predictions))

def compute_accuracy(y_true, y_pred):
    correct = [
        1 if y_true[i] == y_pred[i] else 0
        for i in range(len(y_pred))
    ]
    return 100 * np.sum(correct) / len(y_pred)


def predict_categorical_field_from_neighbours(graph, species_taxid, trait, control_for_phylum=True):
    current_node = graph.nodes[species_taxid]
    
    if pd.isnull(current_node[trait]):
        return None, None, None
    
    current_phylum = current_node['phylum']
    values = collections.defaultdict(int)
    for key in graph.neighbors(species_taxid):
        node = graph.nodes[key]
        
        if pd.isnull(node[trait]):
            continue
        
        if control_for_phylum and node['phylum'] != current_phylum:
            v = node[trait]
            values[v] += 1
        elif not control_for_phylum:
            v = node[trait]
            values[v] += 1
            
    if len(values) == 0:
        prediction = None
        sorted_values = []
    else:
        sorted_values = sorted(values.items(), key=lambda t: t[1], reverse=True)
        prediction = sorted_values[0][0]

    return current_node[trait], prediction, sorted_values

def compute_accuracy_from_neighbour_predictions(graph, species_taxids, trait, control_for_phylum=True):
    actuals = []
    predictions = []
    for i, species_taxid in enumerate(species_taxids):
        actual, prediction, _ = predict_categorical_field_from_neighbours(
            graph, 
            species_taxid, 
            trait, 
            control_for_phylum,
        )
        
        if actual is None or prediction is None:
            continue
        
        actuals.append(actual)
        predictions.append(prediction)
        
    return compute_accuracy(actuals, predictions)

def protein_description_word_count(assembly_accession, below_threshold=True, count_type='document'):
    if count_type not in ('document', 'term'):
        raise ValueError('count type should be one of document or term')
        
    cds_codon_bias_folder = os.path.join(os.getcwd(), 'data/cds_codon_bias')
    if below_threshold:
        cds_codon_bias_folder = os.path.join(cds_codon_bias_folder, 'below_threshold')
    else:
        cds_codon_bias_folder = os.path.join(cds_codon_bias_folder, 'all')
        
    cds_codon_bias_path = os.path.join(cds_codon_bias_folder, f'{assembly_accession}_codon_bias.csv')
    
    black_list = [
        '',
        'protein',
        'enzyme',
        'subunit',
        'domain',
        'domain-containing',
        'hypothetical',
        'putative',
        'probable',
        'conserved',
        'factor',
        'large',
        'small',
        'family',
        'superfamily',
        'outer',
        'cell',
        'i',
        'ii',
        'iii',
        'iv',
        'v',
        'vi',
        'alpha',
        'beta',
        'gamma',
        'type',
        'and',
        'or',
        'of',
        'class',
        'assembly',
        'acid',
        'component',
        'response',
        '2x',
        '1-2',
    ]
    black_list += [c for c in string.ascii_lowercase]
    black_list += [c for c in string.digits]
    black_list = set(black_list)
    
    df = pd.read_csv(cds_codon_bias_path)
    descriptions = df['description'].values
    
    counts = collections.defaultdict(int)
    for description in descriptions:
        desc = str(description).replace(',', '').replace('.', '').replace('(', '').replace(')', '').strip()
        parts_1 = desc.split(' ')
        parts = []
        for p in parts_1:
            parts  += [v.lower() for v in p.split('/')]
            
        for word in parts:
            if word in black_list:
                continue
                
            if count_type == 'document':
                counts[word] = 1
            else:
                counts[word] += 1
    
    return dict(counts)


def compute_word_count(assembly_accessions, below_threshold=True, count_type='document', sorted_output=True):
    if count_type not in ('document', 'term'):
        raise ValueError('count type should be one of document or term')
    
    total_counts = {}
    for i, assembly_accession in enumerate(assembly_accessions):
        if (i+1) % 200 == 0:
            print(f'{i+1} / {len(assembly_accessions)}')
            
        counts = protein_description_word_count(assembly_accession, below_threshold, count_type)
        
        for k, v in counts.items():
            total_counts[k] = total_counts.get(k, 0) + counts[k]
    
    if sorted_output:
        word_counts = [
            (
                k, 
                np.round(100 * v / len(assembly_accessions), 2)
            ) 

            for k, v in total_counts.items()
        ]

        return sorted(word_counts, key=lambda t: t[1], reverse=True)
    else:
        return total_counts

def compute_word_scores(assembly_accessions, min_term_freq=1):
    top_term_freq = compute_word_count(
        assembly_accessions,
        below_threshold=True,
        count_type='term',
        sorted_output=False,
    )
    all_term_freqs = compute_word_count(
        assembly_accessions,
        below_threshold=False,
        count_type='term',
        sorted_output=False,
    )
    all_words = sorted(all_term_freqs.keys())
    
    word_scores = []
    for word in all_words:
        term_freq = top_term_freq.get(word, 0)
        all_term_freq = all_term_freqs[word]
        
        term_freq_norm = term_freq / len(assembly_accessions)
        all_term_freq_norm = all_term_freq / len(assembly_accessions)
        
        if term_freq_norm < min_term_freq:
            continue
            
        score = term_freq / all_term_freqs[word]
        
        word_scores.append((
            word, 
            score, 
            np.round(term_freq_norm, 2), 
            np.round(all_term_freq_norm, 2),
        ))
        
    return sorted(word_scores, key=lambda t: t[1], reverse=True)

def local_word_scores(graph, species_taxid, control_for_phylum=True, include_neighbours=True):
    current_node = graph.nodes[species_taxid]

    current_word_scores = compute_word_scores(current_node['assembly_accessions'])

    score_dicts = [{t[0]: t for t in current_word_scores}]
    word_set = set([t[0] for t in current_word_scores])
    
    neighbours = list(graph.neighbors(species_taxid))
    
    if include_neighbours:
        for key in neighbours:
            node = graph.nodes[key]

            if control_for_phylum and node['phylum'] == current_node['phylum']:
                continue

            word_scores = compute_word_scores(node['assembly_accessions'])

            word_set &= set([t[0] for t in word_scores])
            score_dicts.append({t[0]: t for t in word_scores})

    final_word_scores = []
    for word in sorted(word_set):
        scores, term_freqs, all_term_freqs = [], [], []
        for score_dict in score_dicts:
            _, score, term_freq, all_term_freq = score_dict[word]
            scores.append(score)
            term_freqs.append(term_freq)
            all_term_freqs.append(all_term_freq)
            
        final_word_scores.append((
            word, 
            np.round(np.mean(score), 3),
            np.round(np.mean(term_freqs), 2), 
            np.round(np.mean(all_term_freqs), 2),
        ))
        
    return sorted(final_word_scores, key=lambda t: t[1], reverse=True)

def retrieve_orginal_proteins(graph, species_taxid, include_neighbours=True):
    cds_codon_bias_folder = os.path.join(os.getcwd(), 'data/cds_codon_bias/below_threshold')
    
    species_taxids = [species_taxid] 
    if include_neighbours:
        species_taxids += [key for key in graph.neighbors(species_taxid)]
    
    word_scores = local_word_scores(graph, species_taxid, include_neighbours=include_neighbours)
    words = [t[0] for t in word_scores]
        
    protein_dict = collections.defaultdict(list)
    for species_taxid in species_taxids:
        node = graph.nodes[species_taxid]
        assembly_accessions = node['assembly_accessions']
        
        for assembly_accession in assembly_accessions:
            cds_codon_bias_path = os.path.join(cds_codon_bias_folder, f'{assembly_accession}_codon_bias.csv')
            df = pd.read_csv(cds_codon_bias_path)
            descriptions = df['description'].values
            distances = df['distance'].values
            
            for word in words:
                for i, protein_description in enumerate(descriptions):
                    protein_desc = str(protein_description)
                    protein_desc_key = protein_desc.replace(',', '').replace(': ', ' ').replace(':', ' ').lower()
                    if word.lower() in protein_desc_key:
                        protein_dict[protein_desc_key].append((protein_desc, distances[i]))
                        
    output_proteins = [
        (protein_dict[protein_desc][0][0], np.mean([t[1] for t in protein_dict[protein_desc]]))
        for protein_desc in protein_dict.keys()
    ]
                        
    return sorted(output_proteins, key=lambda t: t[1])


def retrieve_orginal_proteins_global(graph, species_taxids):
    cds_codon_bias_folder = os.path.join(os.getcwd(), 'data/cds_codon_bias/below_threshold')

    protein_dict = collections.defaultdict(list)
    seen_species = set()
    for s_ix, species_taxid in enumerate(species_taxids):
        if (s_ix + 1) % 10 == 0:
            print(f'{s_ix + 1:,} / {len(species_taxids):,}')
        
        if species_taxid in seen_species:
            continue
        
        node = graph.nodes[species_taxid]
        assembly_accessions = node['assembly_accessions']
        
        neighbors = list(graph.neighbors(species_taxid))
        
        seen_species.add(species_taxid)
        seen_species |= set(neighbors)
        
        word_scores = local_word_scores(graph, species_taxid)
        words = [t[0] for t in word_scores]
        
        for assembly_accession in assembly_accessions:
            cds_codon_bias_path = os.path.join(cds_codon_bias_folder, f'{assembly_accession}_codon_bias.csv')
            df = pd.read_csv(cds_codon_bias_path)
            descriptions = df['description'].values
            distances = df['distance'].values
            
            for word in words:
                for i, protein_description in enumerate(descriptions):
                    protein_desc = str(protein_description)
                    protein_desc_key = protein_desc.replace(',', '').replace(': ', ' ').replace(':', ' ').lower()
                    if word.lower() in protein_desc_key:
                        protein_dict[protein_desc_key].append((protein_desc, distances[i]))
                        
    output_proteins = [
        (protein_dict[protein_desc][0][0], np.mean([t[1] for t in protein_dict[protein_desc]]))
        for protein_desc in protein_dict.keys()
    ]
                        
    return sorted(output_proteins, key=lambda t: t[1])

#for v in [t[0] for t in original_proteins_2336 if t[1] < 0.1][:51]:
#    print(v)

ratio_columns = [c for c in species_codon_df.columns if c.endswith('_ratio')]

def load_trnas(engine, assembly_accession):
    trna_query = """
    select * from sequences where sequence_type = 'tRNA' and assembly_accession = ?
    """
    trna_df = pd.read_sql(trna_query, engine, params=(assembly_accession,))
    metadata = [json.loads(v) if v is not None else {} for v in trna_df['metadata_json'].values]
    trna_df['codon'] = [v.get('codon') for v in metadata]
    trna_df['anticodon'] = [v.get('anticodon') for v in metadata]
    trna_df['amino_acid'] = [v.get('amino_acid') for v in metadata]
    return trna_df

adaptation_ix = trna_ai['adaptation_index'].values

matching_df = trna_ai[
    trna_ai['protein'].isin(np.unique([t[0] for t in original_proteins_562]))
]
adaptation_ix_top_codon = matching_df['adaptation_index'].values

def compute_adaptation_mean_diff(assembly_accessions):
    cached_proteins = {}
    
    diff = []
    for i, assembly_accession in enumerate(assembly_accessions):
        if (i+1) % 100 == 0:
            print(f'{i+1:,} / {len(assembly_accessions):,}')
        
        trna_ai_path = os.path.join(os.getcwd(), f'data/trn_adaptation_index/{assembly_accession}_tai.csv')
        trna_ai = pd.read_csv(trna_ai_path)
        
        species_taxid = trna_ai['species_taxid'].iloc[0]
    
        if species_taxid in cached_proteins:
            proteins = cached_proteins[species_taxid]
        else:
            proteins = retrieve_orginal_proteins(
                graph,
                species_taxid=species_taxid,
                include_neighbours=False,
            )
            cached_proteins[species_taxid] = proteins
            
        mean_1 = trna_ai['adaptation_index'].mean()
        mean_2 = trna_ai[
            trna_ai['protein'].isin(np.unique([t[0] for t in proteins]))
        ]['adaptation_index'].mean()
        
        if pd.isnull(mean_1) or pd.isnull(mean_2):
            print(species_taxid, assembly_accession, mean_1, mean_2)
            continue
        
        diff.append(mean_2 - mean_1)
        
    return np.mean(diff), np.std(diff), np.min(diff), np.max(diff)

q = """
select assembly_accession, species_taxid from assembly_source
"""
response_df = pd.read_sql(q, engine)
assembly_accessions = response_df['assembly_accession'].values

def compute_mean_adaptation_per_species(assembly_accessions):
    cached_proteins = {}
    
    res = collections.defaultdict(list)
    for i, assembly_accession in enumerate(assembly_accessions):
        if (i+1) % 100 == 0:
            print(f'{i+1:,} / {len(assembly_accessions):,}')
        
        trna_ai_path = os.path.join(os.getcwd(), f'data/trn_adaptation_index/{assembly_accession}_tai.csv')
        trna_ai = pd.read_csv(trna_ai_path)
        
        species_taxid = trna_ai['species_taxid'].iloc[0]
    
        if species_taxid in cached_proteins:
            proteins = cached_proteins[species_taxid]
        else:
            proteins = retrieve_orginal_proteins(
                graph,
                species_taxid=species_taxid,
                include_neighbours=False,
            )
            cached_proteins[species_taxid] = proteins
            
        mean = trna_ai['adaptation_index'].mean()
        res[species_taxid].append(mean)
        
    output = {}
    for key in sorted(res.keys()):
        output[key] = np.mean([m for m in res[key] if not pd.isnull(m)])
        
    return output

def compute_distributions_of_length(assembly_accessions):
    cds_codon_bias_folder = os.path.join(os.getcwd(), 'data/cds_codon_bias/')
    
    lengths_all = []
    lengths_below = []
    for i, assembly_accession in enumerate(assembly_accessions):
        if (i+1) % 200 == 0:
            print(f'{i+1} / {len(assembly_accessions)}')
        
        path_all = os.path.join(cds_codon_bias_folder, f'all/{assembly_accession}_codon_bias.csv')
        path_below = os.path.join(cds_codon_bias_folder, f'below_threshold/{assembly_accession}_codon_bias.csv')
        
        df_all = pd.read_csv(path_all)
        df_below = pd.read_csv(path_below)
        
        lengths_all.extend(df_all['length'].values.tolist())
        lengths_below.extend(df_below['length'].values.tolist())
        
    return np.array(lengths_all), np.array(lengths_below)

def plot_distributions(lengths_all, lengths_below):
    palette = sns.color_palette()
    _, axes = plt.subplots(1, 2, figsize=(20, 6))
    
    ax1 = axes[0]
    ax2 = axes[1]
    
    ax1.hist(lengths_all, bins=100, log=True)
    ax1.set_title('Length distribution (all)')
    xlim = ax1.get_xlim()
    
    ax2.hist(lengths_below, bins=100, color=palette[1], log=True)
    ax2.set_title('Length distribution (below threshold)')
    ax2.set_xlim(xlim)
    
    return ax1, ax2

distance_matrix_path = os.path.join(os.getcwd(), 'data/distance_matrix.npy')
distance_matrix = np.load(distance_matrix_path, allow_pickle=True)

def make_uniform_bias(codon_to_aa_table, aa_to_codons, ratio_columns):
    uniform_bias = np.zeros((len(ratio_columns),))
    for i, col in enumerate(ratio_columns):
        codon = col.replace('_ratio', '')
        aa = codon_to_aa_table[codon]
        synonyms = aa_to_codons[aa]
        
        uniform_bias[i] = 1 / len(synonyms)
        
    return uniform_bias

def compute_distance_to_uniform_bias(species_codon_df, uniform_bias, ratio_columns):
    X = species_codon_df[ratio_columns].values
    return compute_codon_bias_distance(X, uniform_bias)

distance_to_uniform = compute_distance_to_uniform_bias(species_codon_df, uniform_bias, ratio_columns)
species_codon_df['uniform_dist'] = distance_to_uniform