import pandas as pd
import argparse,json
from collections import defaultdict
from datetime import datetime
import numpy as np
import matplotlib as mpl
mpl.rcParams['axes.formatter.useoffset'] = False

import matplotlib.pyplot as plt
import seaborn as sns
from root_to_tip import filter_and_transform, get_clade_gts

if __name__=="__main__":
    parser = argparse.ArgumentParser(
        description="remove time info",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter
    )

    parser.add_argument('--metadata', type=str, nargs='+', required=True, help="input data")
    parser.add_argument('--clade-gts', type=str, required=True, help="input data")
    parser.add_argument('--clade', type=str, required=True, help="input data")
    parser.add_argument('--sub-clades', type=str, required=True, help="input data")
    parser.add_argument('--min-date', type=float, help="input data")
    parser.add_argument('--output-plot', type=str, help="plot file")
    #parser.add_argument('--output-json', type=str, help="rate file")
    args = parser.parse_args()

    clade_gt = get_clade_gts(args.clade_gts, args.sub_clades)

    d = pd.concat([pd.read_csv(x, sep='\t').fillna('') for x in args.metadata])
    filtered_data, _ = filter_and_transform(d, clade_gt, min_date=args.min_date, max_date=args.min_date + 0.3, completeness=0, swap_root=args.clade_gts=='19B+')
    #filtered_data=filtered_data.loc[filtered_data.country!='China']

    intra_subs_dis = filtered_data["divergence"].value_counts().sort_index()
    intra_aaSubs_dis = filtered_data["aaDivergence"].value_counts().sort_index()
    intra_geno = filtered_data["intra_substitutions_str"].value_counts()


    ls = ['-', '--', '-.', ':']
    plt.figure()
    dates = sorted(filtered_data.loc[:,"numdate"])
    plt.plot(dates, (np.arange(len(dates))+1), c='k', lw=3, alpha=0.3, label="all")
    ls_counter = defaultdict(int)
    for x,i in intra_geno.items():
        nmuts = len(x.split(',')) if x else 0
        ind = filtered_data["intra_substitutions_str"]==x
        dates = sorted(filtered_data.loc[ind,"numdate"])
        factor = 3 if 'C8782T' in x else 1
        if dates[0]<args.min_date+0.2:
        #if (nmuts<2 or i>len(filtered_data)/100) and ls_counter[nmuts]<10:
            if nmuts>5:
                continue
            plt.plot(dates, factor*(np.arange(i)+1), c=f'C{nmuts}', lw=2 if nmuts else 3,
                            ls=ls[ls_counter[nmuts]%len(ls)], marker='o' if len(dates)<3 else '',
                            label=f"gt={x if x else 'founder'}"[:30])
            ls_counter[nmuts]+=1
    plt.ylim(0.5,1000)
    plt.xlim(args.min_date,args.min_date+0.3)
    plt.yscale('log')
    plt.legend()

    if args.output_plot:
        plt.savefig(args.output_plot)
    else:
        plt.show()

import argparse
import numpy as np
import pandas as pd
from scipy.stats import linregress
import matplotlib as mpl
mpl.rcParams['axes.formatter.useoffset'] = False
from root_to_tip import add_panel_label

import matplotlib.pyplot as plt
import seaborn as sns


if __name__=="__main__":
    parser = argparse.ArgumentParser(
        description="remove time info",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter
    )
    fs=14
    parser.add_argument('--rate-table', type=str, required=True, help="input data")
    parser.add_argument('--output-plot', type=str, help="plot file")
    parser.add_argument('--output-plot-rates', type=str, help="plot file")
    parser.add_argument('--output-plot-rates-genes', type=str, help="plot file")
    args = parser.parse_args()

    L = 29903
    Laa = 9716  # all CDS other than ORF9)

    rates = pd.read_csv(args.rate_table, sep='\t', index_col='clade')
    inter_clade_rates = {}
    fig, axs = plt.subplots(2,2, figsize=(12,12))
    xticks = [2020,2020.5,2021,2021.5,2022,2022.5]
    for ax, mut_type, ax_label, panel in zip(axs.flatten(),['nuc', 'aa', 'syn'],
                        ['total divergence', 'amino acid divergence', 'synonymous divergence'], ['A', 'B', 'C']):
        ax.set_ylabel(ax_label, fontsize=fs)
        add_panel_label(ax, panel, fs=fs*1.8)
        inter_clade = []
        ci=0
        for clade, row in rates.iterrows():
            dt = np.linspace(0, 0.75,2)
            t = dt + row[f'{mut_type}_origin']
            slope = row[f'{mut_type}_rate']
            ls = '-' if ci<10 else '--'
            m = 'o' if ci<10 else 's'
            ax.plot(t, dt*slope + row[f'{mut_type}_div'], label=clade, c=f"C{ci%10}", ls=ls)
            ax.scatter([[row[f'{mut_type}_origin']]], [[row[f'{mut_type}_div']]], c=f"C{ci%10}", marker=m)
            if row[f'{mut_type}_origin']>2019.7 and row[f'{mut_type}_origin']<2022.7:
                inter_clade.append([row[f'{mut_type}_origin'], row[f'{mut_type}_div']])
            ci += 1
        inter_clade = np.array(inter_clade)
        reg = linregress(inter_clade[:,0], inter_clade[:,1])
        inter_clade_rates[mut_type] = reg.slope

        x = np.linspace(2019.8, 2022.5, 21)
        y = reg.slope*x + reg.intercept
        std_dev = np.sqrt(np.maximum(0,reg.slope*x + reg.intercept))
        ax.plot(x, y, c='k', lw=3, alpha=0.5)
        ax.fill_between(x, y+std_dev, np.maximum(0, y-std_dev), fc='k', lw=3, alpha=0.1, ec=None)
        ax.set_xlim(x.min(), x.max())
        ax.set_ylim(0)
        if mut_type=='nuc':
            ax.text( 0.5, 0.05,f"overall rate:\n{reg.slope:1.1f}/year\n{reg.slope/L:1.1e}/year/site",
                    fontsize=fs, transform=ax.transAxes)
        else:
            ax.text( 0.5, 0.05,f"overall rate:\n{reg.slope:1.1f}/year\n{reg.slope/Laa:1.1e}/year/codon",
                   fontsize=fs, transform=ax.transAxes)
        if mut_type == 'aa': ax.legend(ncol=2)
        ax.set_xticks(xticks, [str(x) for x in xticks], fontsize=fs)

    for i, (mut_type, rate) in enumerate(inter_clade_rates.items()):
        axs[-1,-1].fill_between([i-0.5, i+0.5], [35,35], facecolor='k', alpha=0.075*(2+i%2))
        axs[-1,-1].plot([i-0.4, i+0.4], [rate, rate], lw=3, c='k', alpha=0.5)
        clade_rates = rates[f"{mut_type}_rate"]
        x_offset = i + np.linspace(-.35, 0.35,len(clade_rates))
        axs[-1,-1].scatter(x_offset[:10],
                           clade_rates[:10], c=[f"C{ci%10}" for ci in range(10)],
                           marker='o')
        axs[-1,-1].scatter(x_offset[10:],
                           clade_rates[10:], c=[f"C{ci%10}" for ci in range(len(clade_rates) - 10)],
                           marker='s')
    axs[-1,-1].set_ylabel("substitutions per year", fontsize=fs)
    axs[-1,-1].set_xticks([0,1,2], ['total', 'amino acid', 'synonymous'], rotation=20,
                          ha='center', fontsize=fs)
    axs[-1,-1].set_ylim(0)
    add_panel_label(axs[-1,-1], 'D', fs=fs*1.8)

    if args.output_plot:
        plt.savefig(args.output_plot)
    else:
        plt.show()

    plt.figure()
    plt.plot(rates["nuc_origin"], rates["aa_rate"], 'o', label='amino acid rate')
    plt.plot(rates["nuc_origin"], rates["syn_rate"], 'o', label='synonymous rate')
    plt.plot(rates["nuc_origin"], np.ones_like(rates['nuc_origin'])*inter_clade_rates["aa"], label='inter-clade amino acid rate', c=f"C{0}", lw=3)
    plt.plot(rates["nuc_origin"], np.ones_like(rates['nuc_origin'])*inter_clade_rates["syn"], label='inter-clade synonymous rate', c=f"C{1}", lw=3)
    plt.ylabel('rate estimate [1/y]')
    plt.ylim(0)
    plt.legend()

    if args.output_plot_rates:
        plt.savefig(args.output_plot_rates)
    else:
        plt.show()

    plt.figure()
    plt.plot(rates["aa_rate"], 'o-', label='Overall amino acid rate')
    plt.plot(rates["spike_rate"], 's-', label='spike protein')
    plt.plot(rates["orf1a_rate"], 'd-', label='ORF1a')
    plt.plot(rates["orf1b_rate"], 'd-', label='ORF1b')
    plt.plot(rates["enm_rate"], 'd-', label='E,M,N')
    plt.plot(rates["aa_rate"] - rates["spike_rate"] - rates["orf1a_rate"] - rates["orf1b_rate"]- rates["enm_rate"],
             'v-', label='other ORFs')
    plt.ylabel('rate estimate [subs/y]')
    plt.legend()
    plt.xticks(range(len(rates)), rates.index, rotation=60, ha='right')
    plt.tight_layout()

    if args.output_plot_rates_genes:
        plt.savefig(args.output_plot_rates_genes)
    else:
        plt.show()

import argparse,json
from unittest.mock import NonCallableMagicMock
import pandas as pd
import numpy as np
import matplotlib as mpl
from Bio import SeqIO,Seq
mpl.rcParams['axes.formatter.useoffset'] = False
from collections import defaultdict
from scipy.stats import scoreatpercentile

import matplotlib.pyplot as plt
import seaborn as sns

columns = ['seqName', 'Nextclade_pango', "privateNucMutations.unlabeledSubstitutions", 'qc.overallStatus', 'privateNucMutations.reversionSubstitutions']

def get_sequence(mods, root_seq):
    seq = root_seq.copy()
    for mut in mods['nuc']:
        a,pos,d = mut[0], int(mut[1:-1])-1, mut[-1]
        if a!=seq[pos]:
            print(seq[pos], mut)
        seq[pos]=d
    return seq

def translate(codon):
    try:
        return Seq.translate("".join(codon))
    except:
        return 'X'

if __name__=="__main__":
    parser = argparse.ArgumentParser(
        description="remove time info",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter
    )
    parser.add_argument('--metadata', type=str, required=True, help="input data")
    parser.add_argument('--reference', type=str, required=True, help="input data")
    parser.add_argument('--pango-gts', type=str, required=True, help="input data")
    parser.add_argument('--rare-cutoff', type=int, default=2, help="minimal number of occurrences to count mutations")
    parser.add_argument('--output-fitness', type=str, required=True, help="fitness table")
    parser.add_argument('--output-events', type=str, required=True, help="events table")
    parser.add_argument('--output-mutations', type=str, required=True, help="fitness table")
    args = parser.parse_args()

    ## load reference
    ref = SeqIO.read(args.reference, 'genbank')
    ref_array = np.array(ref.seq)
    base_content = {x:ref.seq.count(x) for x in 'ACGT'}

    ## make a map of codon positions, ignore orf9b (overlaps N)
    codon_pos = np.zeros(len(ref))
    map_to_gene = {}
    gene_position = {}
    gene_length = {}
    for feat in ref.features:
        if feat.type=='CDS':
            gene_name = feat.qualifiers['gene'][0]
            gene_position[gene_name] = feat.location
            gene_length[gene_name] = (feat.location.end - feat.location.start)//3
            if gene_name!='ORF9b':
                for gpos, pos in enumerate(feat.location):
                    codon_pos[pos] = (gpos%3)+1
                    map_to_gene[pos]= gene_name


    # load pango genotypes
    with open(args.pango_gts) as fh:
        pango_gts = json.load(fh)

    # load and filter metadata
    d = pd.read_csv(args.metadata, sep='\t', usecols=columns).fillna('')
    d = d.loc[d["qc.overallStatus"]=='good',:]

    # glob all rare mutations by pango lineage
    mutation_counter = defaultdict(lambda: defaultdict(int))
    pango_counter = defaultdict(int)
    for r, row in d.iterrows():
        pango = row.Nextclade_pango
        if pango[0]=='X':
            continue

        pango_counter[pango] += 1
        muts = row["privateNucMutations.unlabeledSubstitutions"].split(',')
        rev_muts = row["privateNucMutations.reversionSubstitutions"].split(',')
        if len(muts) + len(rev_muts):
            for m in muts + rev_muts:
                if m:
                    mutation_counter[pango][m] += 1

    lineage_size_cutoff = 100
    big_lineages = [pango for pango in pango_counter if pango_counter[pango]>lineage_size_cutoff]
    nlin = len(big_lineages)
    tmp_number_of_pango_muts = defaultdict(int)
    for pango, muts in pango_gts.items():
        if pango_counter[pango]>lineage_size_cutoff:
            for m in muts['nuc']:
                tmp_number_of_pango_muts[int(m[1:-1])-1] += 1
    number_of_pango_muts = np.array([tmp_number_of_pango_muts.get(pos,0)/nlin for pos in range(len(ref))])


    # for each pango lineage, count mutations by position, as well as synoymous and non-synonyomous by codon
    position_counter = defaultdict(lambda: defaultdict(list))
    position_freq = defaultdict(lambda: defaultdict(list))
    syn =    defaultdict(lambda: defaultdict(lambda: defaultdict(list)))
    nonsyn = defaultdict(lambda: defaultdict(lambda: defaultdict(list)))
    for pango, pmuts in mutation_counter.items():
        if pango not in pango_gts:
            print("missing pango", pango)
            continue
        seq = get_sequence(pango_gts[pango], ref_array)
        # translate each codon in pango reference sequence
        for mut,c in pmuts.items():
            a, pos, d = mut[0], int(mut[1:-1])-1, mut[-1]
            if a=='-':
                continue
            position_counter[pos][(a,d)].append(c)
            if pango_counter[pango]>lineage_size_cutoff and c>1:
                position_freq[pos][(a,d)].append(c/pango_counter[pango])

            if pos in map_to_gene:
                gene = map_to_gene[pos]
                pos_in_gene = pos - gene_position[gene].start
                frame = pos_in_gene%3
                codon = seq[pos-frame:pos-frame+3]
                codon_number = pos_in_gene//3
                alt = codon.copy()
                alt[frame] = d
                aa = translate(codon)
                alt_aa = translate(alt)
                if aa==alt_aa and aa!='X':
                    syn[gene][codon_number][a].append(c)
                elif 'X' not in [aa,alt_aa]:
                    nonsyn[gene][codon_number][a].append(c)

    sorted_transitions = []
    for a in 'ACGT':
        for d in 'ACGT':
            if a!=d:
                sorted_transitions.append((a,d))

    # determine the total number of events and counts at each position
    def total_len_filtered_lists(l, cutoff, rates=None):
        if rates is None:
            rates = {}
        return sum([len([y for y in x if y>=cutoff])/rates.get(a,1.0) for (a,d),x in l.items()])

    all_events = []
    for pos in range(len(ref)):
        tmp = []
        for a, d in sorted_transitions:
            tmp.append(total_len_filtered_lists({(a,d): position_counter[pos][(a,d)]}, args.rare_cutoff))
        all_events.append(tmp)

    # sort all mutations by type of mutation
    events_by_transition = defaultdict(int)
    for m in position_counter.values():
        for t in m:
            events_by_transition[t]+=len(m[t])
    # determine the mutations rate for each type, as well as the rate out of a nuc
    transition_rate = dict()
    away_rate = defaultdict(float)
    total_rate = np.sum(list(events_by_transition.values()))
    for a,d in events_by_transition:
        if a in base_content:
            transition_rate[(a,d)] = events_by_transition[(a,d)]/base_content[a]/total_rate*len(ref)
            away_rate[a] += transition_rate[(a,d)]


    total_events = np.array([total_len_filtered_lists(position_counter[pos], args.rare_cutoff)
                            if pos in position_counter else 0
                            for pos in range(len(ref))])

    total_events_weighted = np.array([total_len_filtered_lists(position_counter[pos], args.rare_cutoff, away_rate)
                            if pos in position_counter else 0
                            for pos in range(len(ref))])

    total_counts = np.array([sum([sum(x) for x in position_counter[pos].values()])
                             if pos in position_counter else 0
                             for pos in range(len(ref))])

    # average mutation rate used to scale rates
    total_events_rescaled = np.array([c/away_rate[nuc] for c,nuc in zip(total_events, ref.seq)])
    total_events_rescaled /= np.median(total_events_rescaled)
    total_events_weighted /= np.median(total_events_weighted)

    with open(args.output_events, 'w') as fh:
        fh.write('\t'.join(['position', 'ref_state', 'lineage_fraction_with_changes', 'gene', 'codon', 'pos_in_codon']
                            + [f"{a}->{d}" for a,d in sorted_transitions]) + '\n')
        for pos, nuc in enumerate(ref_array):
            gene = map_to_gene.get(pos,"")
            data = "\t".join([f"{all_events[pos][i]}" for i in range(len(sorted_transitions))])
            if gene:
                gene_pos = pos - gene_position[gene].start
                codon = gene_pos//3 + 1
                cp = gene_pos%3 + 1
                fh.write(f'{pos+1}\t{nuc}\t{number_of_pango_muts[pos]:1.3f}\t{gene}\t{codon}\t{cp}\t' + data + "\n")
            else:
                fh.write(f'{pos+1}\t{nuc}\t{number_of_pango_muts[pos]:1.3f}\t\t\t\t' + data + "\n")

    with open(args.output_fitness, 'w') as fh:
        fh.write('\t'.join(['position', 'ref_state', 'lineage_fraction_with_changes', 'gene', 'codon', 'pos_in_codon', 'total_count', 'total_events', 'tolerance']) + '\n')
        for pos, nuc in enumerate(ref_array):
            gene = map_to_gene.get(pos,"")
            if gene:
                gene_pos = pos - gene_position[gene].start
                codon = gene_pos//3 + 1
                cp = gene_pos%3 + 1
                fh.write(f'{pos+1}\t{nuc}\t{number_of_pango_muts[pos]:1.3f}\t{gene}\t{codon}\t{cp}\t{total_counts[pos]}\t{total_events[pos]}\t{total_events_weighted[pos]:1.3f}\n')
            else:
                fh.write(f'{pos+1}\t{nuc}\t{number_of_pango_muts[pos]:1.3f}\t\t\t\t{total_counts[pos]}\t{total_events[pos]}\t{total_events_weighted[pos]:1.3f}\n')

    with open(args.output_mutations, 'w') as fh:
        fh.write('\t'.join(['mutation', 'raw_counts', 'origin_sites', 'scaled_rate']) + '\n')
        for a in 'ACTG':
            for d in 'ACGT':
                if a==d: continue
                key = (a,d)
                fh.write(f"{a}->{d}\t{events_by_transition[key]}\t{base_content[a]}\t{transition_rate[key]:1.3f}\n")


    # # calculate synonymous and non-synonymous distributions
    # total_events_by_type = {}
    # for mut_counts, label in [(syn, 'syn'), (nonsyn, 'nonsyn')]:
    #     total_events_by_type[label] = {}
    #     for gene in mut_counts:
    #         total_events_by_type[label][gene] = [np.sum([len([x for x in mut_counts[gene][pos][a] if x>cutoff])/away_rate[a]
    #                                                  for a in mut_counts[gene][pos]])
    #                                                for pos in range(gene_length[gene])]


    # def calc_fitness_cost(freqs, transition):
    #     mu = 0.0004/50
    #     mut_rate = transition_rate[transition]*mu
    #     avg_freq = np.sum(freqs)/nlin
    #     return mut_rate/(avg_freq + 1e-3/nlin)

    # fitness_cost = []
    # for pos,nuc in enumerate(ref_array):
    #     total_muts = np.sum([len(v) for v in position_freq[pos].values()])
    #     reference_muts = np.sum([len(v) for (a,d), v in position_freq[pos].items() if a==nuc])
    #     if reference_muts>0.9*total_muts:
    #         fitness_cost.append([np.inf if d==nuc
    #                              else calc_fitness_cost(position_freq[pos][(nuc,d)],(nuc,d))
    #                             for d in 'ACGT'])
    #     else:
    #         fitness_cost.append([np.nan, np.nan, np.nan, np.nan])