#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Run future simulations based on IUCN data and status transition rates

Created on Wed Oct 30 20:59:28 2019
@author: Tobias Andermann (tobias.andermann@bioenv.gu.se)
"""

import numpy as np
np.set_printoptions(suppress=True)
import pandas as pd
import matplotlib.pyplot as plt
from matplotlib.backends.backend_pdf import PdfPages
import os
import sys
import iucn_sim.functions as cust_func

def add_arguments(parser):
    parser.add_argument(
        '--input_data',
        required=True,
        help="Path to 'simulation_input_data.pkl' file created by esimate_rates function."
    )
    parser.add_argument(
        '--outdir',
        required=True,
        help="Provide path to outdir where results will be saved."
    )
    parser.add_argument(
        '--n_years',
        default=100,
        help="How many years to simulate into the future."
    )
    parser.add_argument(
        '--n_sim',
        default=10000,
        help="How many simulation replicates to run. At least 10,000 simulations are recommended for accurate rate estimation (default). If the number of simulation replicates exceeds the number of available transition rate estimates (produced by the 'transition_rates' function), these rates will be randomely resampled for the remaining simulations."
    )
    parser.add_argument(
        '--status_change',
        default=1,
        help="Model IUCN status changes in future simulations. 0=off, 1=on (default=1)."
    )
    parser.add_argument(
        '--conservation_increase_factor',
        default=1,
        help="The transition rates leading to improvements in IUCN conservation status are multiplied by this factor."
    )
    parser.add_argument(
        '--threat_increase_factor',
        default=1,
        help="Opposite of conservation_increase_factor, multiplies the transition rates leading to worsening in IUCN conservation status."
    )
    parser.add_argument(
        '--model_unknown_as_lc',
        default=0,
        help="Model new status for all DD and NE species as LC (best case scenario). 0=off, 1=on (default=0)."
    )
    parser.add_argument(
        '--n_extinct_taxa',
        default=0,
        help="Setting this value will stop the simulations when n taxa have gone extinct. This can be used to simulate the expected time until n extinctions. The value of the --n_years flag in this case will be interpreted as the maximum possible time frame, so set it large enough to cover a realistic time-frame for these extinctions to occur. Set to 0 to disable this function (default=0)."
    )
    parser.add_argument(
        '--extinction_rates',
        default=1,
        help="Estimation of extinction rates from simulation results: 0=off, 1=on (default=1)."
    )
    parser.add_argument(
        '--n_gen',
        default=100000,
        help="Number of generations for MCMC for extinction rate estimation (default=100000)."
    )
    parser.add_argument(
        '--burnin',
        default=1000,
        help="Burn-in for MCMC for extinction rate estimation (default=1000)."
    )
    parser.add_argument(
        '--plot_diversity_trajectory',
        default=1,
        help="Plots the simulated diversity trajectory: 0=off, 1=on (default=1)."
    )
    parser.add_argument(
        '--plot_status_trajectories',
        default=1,
        help="Plots the simulated IUCN status trajectory: 0=off, 1=on (default=0)."
    )
    parser.add_argument(
        '--plot_histograms',
        default=0,
        help="Plots histograms of simulated extinction times for each species: 0=off, 1=on (default=0)."
    )
    parser.add_argument(
        '--plot_posterior',
        default=0,
        help="Plots histograms of posterior rate estimates for each species: 0=off, 1=on (default=0)."
    )
    parser.add_argument(
        '--plot_status_piechart',
        default=1,
        help="Plots pie charts of status distribution: 0=off, 1=on (default=1)."
    )
    parser.add_argument(
        '--seed',
        default=None,
        help="Set random seed for future simulations."
    )

def p_e_year(years,p_e):
    pe_year = 1-(1-p_e)**(1/years)
    return pe_year

def update_multiplier(q,d=1.1):
    u = np.random.uniform(0,1)
    l = 2*np.log(d)
    m = np.exp(l*(u-.5))
    new_q = q * m
    return new_q, np.log(m)

def get_rate_estimate(ext_time_array,max_t,index,species_list,plot_posterior=0,pdf=0,n_gen = 100000,burnin = 1000):
    sys.stdout.write('\rProcessing species: %i/%i '%(index+1,len(species_list)))
    ext_time_array_new = ext_time_array.copy()
    ext_time_array_new[ext_time_array_new!=ext_time_array_new] = max_t
    ext_time_array_new = ext_time_array_new.astype(float)
    w_times = np.sum(ext_time_array_new)
    ext_events = len(ext_time_array_new[ext_time_array_new<max_t])
    post_samples = []
    q = 0.01
    likA = np.log(q)*ext_events -q*w_times
    for i in range(n_gen):
        new_q, hast = update_multiplier(q)
        lik = np.log(new_q)*ext_events -new_q*w_times
        if lik-likA + hast >= np.log(np.random.random()):
            q = new_q
            likA = lik
        if i > burnin and i % 10==0:
            post_samples.append(q)
    mean_value = np.mean(post_samples)
    lower,upper = cust_func.calcHPD(post_samples,0.95)
    if plot_posterior:
        plt.figure()
        plt.hist(post_samples,100)
        plt.xlabel('Extinction rate estimates')
        plt.ylabel('Counts')
        plt.title(species_list[index])
        plt.tight_layout()
        pdf.savefig()
        plt.close()
        #fig.savefig(os.path.join(posterior_plot_dir,'%s.pdf'%species_list[index]),bbox_inches='tight', dpi = 500)
    return [mean_value,lower,upper]

def select_target_species(species,species_list_status,species_list,en_ext_data,cr_ext_data):
    target_species = species
    target_index = species_list_status[species_list_status.species==target_species].index.values[0]
    species_list_status = species_list_status.iloc[target_index,:]
    species_list = np.array([species_list[target_index]])
    en_ext_data = np.array([en_ext_data[target_index]])
    cr_ext_data = np.array([cr_ext_data[target_index]])
    return pd.DataFrame(species_list_status).T,species_list,en_ext_data,cr_ext_data

def get_rate_estimate_posterior(ext_time_array,max_t,index,species_list,n_gen = 100000,burnin = 1000):
    sys.stdout.write('\rProcessing species: %i/%i '%(index+1,len(species_list)))
    ext_time_array_new = ext_time_array.copy()
    ext_time_array_new[ext_time_array_new!=ext_time_array_new] = max_t
    ext_time_array_new = ext_time_array_new.astype(float)
    w_times = np.sum(ext_time_array_new)
    ext_events = len(ext_time_array_new[ext_time_array_new<max_t])
    post_samples = []
    q = 0.01
    likA = np.log(q)*ext_events -q*w_times
    for i in range(n_gen):
        new_q, hast = update_multiplier(q)
        lik = np.log(new_q)*ext_events -new_q*w_times
        if lik-likA + hast >= np.log(np.random.random()):
            q = new_q
            likA = lik
        if i > burnin and i % 10==0:
            post_samples.append(q)
    return post_samples

## test rate estimator
#true_rate = 0.01
#n_sim = 100
#ext_time_array = (np.random.exponential(1./true_rate, n_sim)).astype(int)
#max_t=100
#ext_time_array[ext_time_array>max_t] = max_t
#get_rate_estimate(ext_time_array,max_t)

def main(args):

    # get user input___________________________________________________________
    seed = args.seed
    try:
        random_seed = int(seed)
        print('Simulating with user-set starting seed %i.'%random_seed)

    except:
        random_seed = np.random.randint(999999999)
        print('Simulating with randomely generated starting seed %i.'%random_seed)
    np.random.seed(random_seed)

    infile = args.input_data
    outdir = args.outdir
    n_years = int(args.n_years)
    n_sim = int(args.n_sim)
    n_extinct_taxa = int(args.n_extinct_taxa)
    allow_status_change = int(args.status_change)
    conservation_increase_factor = int(args.conservation_increase_factor)
    threat_increase_factor = int(args.threat_increase_factor)
    extinction_rates = int(args.extinction_rates)
    n_gen = int(args.n_gen)
    burnin = int(args.burnin)
    plot_diversity_trajectory = int(args.plot_diversity_trajectory)
    plot_histograms = int(args.plot_histograms)
    plot_posterior = int(args.plot_posterior)
    model_unknown_as_lc = int(args.model_unknown_as_lc)
    plot_status_trajectories = int(args.plot_status_trajectories)
    plot_status_piechart = int(args.plot_status_piechart)
    if not os.path.exists(outdir):
        os.makedirs(outdir)
    np.savetxt(os.path.join(outdir,'starting_seed.txt'),np.array([random_seed]),fmt='%i')

    input_data = cust_func.load_obj(infile)
    species_input_data, dd_probs = input_data
    species_list = np.array([i[0] for i in species_input_data])
    current_status_list = np.array([i[1] for i in species_input_data])
    q_matrix_list = [i[2] for i in species_input_data]
    n_rates = dd_probs.shape[1]
    #__________________________________________________________________________

    # modify q-matrices, if set by user________________________________________
    final_qmatrix_list = []
    q_matrix_list_copy = np.array(q_matrix_list).copy()
    for q_matrix_list_i in q_matrix_list_copy:
        q_matrix_list_temp = []
        for q_matrix in q_matrix_list_i:
            if conservation_increase_factor != 1:
                indeces_lower_triangle = np.tril_indices(q_matrix.shape[0],-1)
                q_matrix[indeces_lower_triangle] = q_matrix[indeces_lower_triangle] * conservation_increase_factor
                np.fill_diagonal(q_matrix,0)
                np.fill_diagonal(q_matrix, -np.sum(q_matrix,axis=1))
            if threat_increase_factor != 1:
                indeces_upper_triangle = np.triu_indices(q_matrix.shape[0],1)
                q_matrix[indeces_upper_triangle] = q_matrix[indeces_upper_triangle] * threat_increase_factor
                np.fill_diagonal(q_matrix,0)
                np.fill_diagonal(q_matrix, -np.sum(q_matrix,axis=1))
            q_matrix_list_temp.append(q_matrix)
        final_qmatrix_list.append(q_matrix_list_temp)
    # turn into a dict with the n qmatrices for each species
    final_qmatrix_dict = dict(zip(species_list,final_qmatrix_list))
    #__________________________________________________________________________

    # if more n_rep are set than there are q-matrices, resample________________
    if n_sim <= n_rates:
        sample_columns = np.random.choice(np.arange(n_rates),size=n_sim,replace=False)
    # since there are only as many cr and en p(ex) estimates as there are provided GL values, we may have to resample some (but make sure all are present at least once)
    else:
        sample_columns1 = np.random.choice(np.arange(n_rates),size=n_rates,replace=False)
        sample_columns2 = np.random.choice(np.arange(n_rates),size=(n_sim-n_rates),replace=True)
        sample_columns = np.concatenate([sample_columns1,sample_columns2])
    #__________________________________________________________________________

    # run simulations__________________________________________________________
    delta_t = n_years
    model_ne_as_dd = False
    dynamic_qmatrix = True

    if model_ne_as_dd:
        current_status_list[current_status_list=='NE'] = 'DD'
    if model_unknown_as_lc:
        print('\nSetting all DD and NE species to LC.')
        all_lc=True
    else:
        all_lc=False
    if allow_status_change:
        status_change=True
    else:
        print('\nNot simulating future status changes!')
        status_change=False

    print('\nStarting simulations ...')
    diversity_through_time,te_array,status_through_time,time_until_n_extinctions_list = cust_func.run_multi_sim(n_sim,delta_t,species_list,current_status_list,dd_probs,final_qmatrix_dict,sample_columns,outdir,all_lc=all_lc,status_change=status_change,dynamic_qmatrix=dynamic_qmatrix,n_extinct_taxa=n_extinct_taxa)
    # summarize simulation results
    sim_species_list = te_array[:,0].copy()
    ext_date_data = te_array[:,1:].copy()
    extinction_occs = np.array([len(row[~np.isnan(list(row))]) for row in ext_date_data])
    extinction_prob = extinction_occs/ext_date_data.shape[1]
    # produce output file for status distribution through time
    mean_status_through_time = np.mean(status_through_time,axis=2)
    year = np.arange(delta_t+1).astype(int)
    status_df_data = np.round(np.vstack([year,mean_status_through_time])).astype(int)
    status_df = pd.DataFrame(data = status_df_data.T,columns=['year','LC','NT','VU','EN','CR','EX'])
    status_df.to_csv(os.path.join(outdir,'status_distribution_through_time.txt'),sep='\t',index=False)
    np.savetxt(os.path.join(outdir,'simulated_extinctions_array.txt'),status_through_time[-1],fmt='%i')
    pd.DataFrame(data=te_array).to_csv(os.path.join(outdir,'te_all_species.txt'),sep='\t',header=False,index=False)
    #__________________________________________________________________________

    #__________________________________________________________________________
#    if target_species:
#        posterior = get_rate_estimate_posterior(ext_date_data[0],n_years,0,sim_species_list,n_gen=n_gen,burnin=burnin)
#        np.savetxt('/Users/tobias/GitHub/iucn_predictions/doc/figures/Figure_2/figure_data/posterior_samples/%s_gl_no_status_change.txt'%target_species.replace(' ','_'),posterior,fmt='%.8f')
#        print('\nPrinted posterior')
    #__________________________________________________________________________

    if n_extinct_taxa:
        np.savetxt(os.path.join(outdir,'time_until_%i_extinctions.txt'%n_extinct_taxa),time_until_n_extinctions_list,fmt='%.2f')
        fig = plt.figure()
        plt.hist(time_until_n_extinctions_list)
        plt.xlabel('Time in years')
        plt.ylabel('N')
        plt.title('Simulated years until %i extinctions (%i simulations)'%(n_extinct_taxa,n_sim))
        fig.savefig(os.path.join(outdir,'time_until_%i_extinctions.pdf'%n_extinct_taxa),bbox_inches='tight', dpi = 500)

    #__________________________________________________________________________
    if plot_diversity_trajectory:
        # plot diversity trajectory of species list________________________________
        #colors = ["#9a002e","#df4a3d","#fecd5f","#5cd368","#916200"]
        # define time axis
        time_axis = np.array(range(len(diversity_through_time[0])))
        fig = plt.figure()
        y_values = np.mean(diversity_through_time, axis =0)
        plt.plot(time_axis,y_values,color="#b80033", label='accounting for GL')
        # get upper and lower confidence interval boundaries
        min_hpd, max_hpd = np.array([cust_func.calcHPD(i,0.95) for i in diversity_through_time.T]).T
        mean_min_max = np.vstack([y_values,min_hpd,max_hpd])
        np.savetxt(os.path.join(outdir,'future_diversity_trajectory.txt'),mean_min_max,fmt='%.2f')
        plt.fill_between(time_axis, min_hpd, max_hpd,
                 color="#b80033", alpha=0.2)
        #plt.legend()
        plt.ylabel('Total diversity')
        plt.xlabel('Years from present')
        ax = plt.gca()
        ax1 = ax.twinx()
        # Set the limits of the new axis from the original axis limits
        ax1.set_ylim(ax.get_ylim())
        current_diversity = diversity_through_time[0,0]
        plt.yticks([np.mean(diversity_through_time[:,-1])],[int(current_diversity-np.mean(diversity_through_time[:,-1]))])
        #plt.xticks(modified_q_matrix.year[::10],modified_q_matrix.year[::10])
        plt.ylabel('Lost species')
        plt.tight_layout()
        fig.savefig(os.path.join(outdir,'future_diversity_trajectory.pdf'),bbox_inches='tight', dpi = 500)
    #__________________________________________________________________________

    #__________________________________________________________________________
    if plot_status_trajectories:
        # color palette
        colors = ["#227a00","#a5c279","#f3d248","#6956cb","#79262a","#e34349"]
        # define time axis
        time_axis = np.array(range(len(diversity_through_time[0])))
        # plot results
        def plot_mean_and_interval(div,color,label,fig):
            plt.plot(time_axis,np.mean(div,axis=0),color=color,label=label);
            min_hpd, max_hpd = np.array([cust_func.calcHPD(i,0.95) for i in div.T]).T
            plt.fill_between(time_axis, min_hpd, max_hpd, color=color, alpha=0.2);
            return fig
        fig = plt.figure(figsize=(10,10))
        plot_mean_and_interval(status_through_time[0,:,:].T,colors[0],'LC',fig)
        plot_mean_and_interval(status_through_time[1,:,:].T,colors[1],'NT',fig)
        plot_mean_and_interval(status_through_time[2,:,:].T,colors[2],'VU',fig)
        plot_mean_and_interval(status_through_time[3,:,:].T,colors[3],'EN',fig)
        plot_mean_and_interval(status_through_time[4,:,:].T,colors[4],'CR',fig)
        plot_mean_and_interval(status_through_time[5,:,:].T,colors[5],'EX',fig)
        # add title, legend and axis-labels
        plt.legend(loc='best',fancybox=True)
        plt.title('Diversity trajectory IUCN categories - status change') #10x higher conservation
        plt.ylabel('Number species in category')
        plt.xlabel('Years from present')
        ax = plt.gca()
        ax1 = ax.twinx()
        # Set the limits of the new axis from the original axis limits
        ax1.set_ylim(ax.get_ylim())
        # annotate final counts with labels
        right_ticks = [int(np.round(np.mean(status_through_time[i,-1,:]))) for i in range(status_through_time.shape[0])]
        plt.yticks(right_ticks,right_ticks)
        #plt.xticks(modified_q_matrix.year[::10],modified_q_matrix.year[::10])
        plt.tight_layout()
        fig.savefig(os.path.join(outdir,'future_status_trajectory.pdf'),bbox_inches='tight', dpi = 500)
    #__________________________________________________________________________

    #__________________________________________________________________________
    if plot_status_piechart:
        statuses, counts = np.unique(current_status_list,return_counts=True)
        init_status_dict = dict(zip(statuses, counts))
        init_status_dict['EX'] = 0
        iucn_status_code = {0:'LC', 1:'NT', 2:'VU', 3:'EN', 4:'CR', 5:'EX', 6:'DD'}
        status_count_list = []
        for status_id in np.arange(status_through_time.shape[0]+1):
            status = iucn_status_code[status_id]
            if status in init_status_dict.keys():
                pre_dd_modeling_count = init_status_dict[status]
            else:
                pre_dd_modeling_count = 0
            if not status == 'DD':
                present_status_count = int(np.round(np.mean(status_through_time[status_id][0])))
                final_status_count = int(np.round(np.mean(status_through_time[status_id][-1])))
            else:
                present_status_count = 0
                final_status_count = 0
            status_count_list.append([pre_dd_modeling_count,present_status_count,final_status_count])
        status_count_list = np.array(status_count_list).T
        colors = np.array(["#227a00","#a5c279","#f3d248","#6956cb","#79262a","#b80033",'black'])
        labels = np.array(['LC', 'NT', 'VU', 'EN', 'CR', 'EX', 'DD'])
        def func(pct, allvals):
            absolute = int(np.round((pct/100.*np.sum(allvals))))
            return "{:d}".format(absolute)
        fig, axs = plt.subplots(1, 3,figsize=(12,10))
        # status distribution beginning
        wedges, texts, autotexts =axs[1].pie(status_count_list[1][status_count_list[1] >0], colors= colors[status_count_list[1] >0], autopct=lambda pct: func(pct, status_count_list[1][status_count_list[1] >0]), shadow=False,textprops=dict(color="w"))
        # status distribution end
        wedges, texts, autotexts =axs[2].pie(status_count_list[2][status_count_list[2] >0], colors= colors[status_count_list[2] >0], autopct=lambda pct: func(pct, status_count_list[2][status_count_list[2] >0]), shadow=False,textprops=dict(color="w"))
        ext = wedges[-1]
        # status distribution pre-dd
        wedges, texts, autotexts =axs[0].pie(status_count_list[0][status_count_list[0] >0], colors= colors[status_count_list[0] >0], autopct=lambda pct: func(pct, status_count_list[0][status_count_list[0] >0]), shadow=False,textprops=dict(color="w"))
        axs[0].set_title('Current (including DD)')
        axs[1].set_title('Current (DD corrected)')
        axs[2].set_title('Final (%i years)'%delta_t)
        final_labels = list(labels[status_count_list[0] >0]) + ['EX']
        plt.legend(wedges+[ext], final_labels,title="IUCN status\n(N=%i sp.)"%status_count_list[2].sum(),loc="center left",bbox_to_anchor=(1, 0, 0.5, 1))
        fig.savefig(os.path.join(outdir,'status_pie_chart.pdf'),bbox_inches='tight', dpi = 500)
    #__________________________________________________________________________

    #__________________________________________________________________________
    if extinction_rates:
        # calculate some extinction stats
        # estimate extinction rates scaled by year
        print('\nRunning %i MCMCs to estimate species-specific extinction rates from simulation output...'%len(species_list))
        #ext_date_data = ext_date_data[:10,:]
        if plot_posterior:
            with PdfPages(os.path.join(outdir,'posterior_ext_rate_histograms.pdf')) as pdf:
                sampled_rates = np.array([get_rate_estimate(species_values,n_years,i,sim_species_list,plot_posterior=plot_posterior,pdf=pdf,n_gen=n_gen,burnin=burnin) for i,species_values in enumerate(ext_date_data)])
        else:
            sampled_rates = np.array([get_rate_estimate(species_values,n_years,i,sim_species_list,plot_posterior=plot_posterior,pdf=0,n_gen=n_gen,burnin=burnin) for i,species_values in enumerate(ext_date_data)])
        # export extinction stats to file
        column_names = ['species','rate_e_mean','rate_e_lower','rate_e_upper','simulated_p_e_in_%i_years'%delta_t]
        extinction_prob_df = pd.DataFrame(np.array([sim_species_list,sampled_rates[:,0],sampled_rates[:,1],sampled_rates[:,2],extinction_prob]).T,columns=column_names)
    else:
        column_names = ['species','simulated_p_e_in_%i_years'%delta_t]
        extinction_prob_df = pd.DataFrame(np.array([sim_species_list,extinction_prob]).T,columns=column_names)
    extinction_prob_df[column_names[1:]] = extinction_prob_df[column_names[1:]].astype(float)
    extinction_prob_df.to_csv(os.path.join(outdir,'extinction_prob_all_species.txt'),sep='\t',index=False,float_format='%.8f')
    print('\n')
    #__________________________________________________________________________

    #__________________________________________________________________________
    if plot_histograms:
        # plot histograms of extinction times
        with PdfPages(os.path.join(outdir,'extinction_time_histograms.pdf')) as pdf:
            for i,species in enumerate(te_array[:,0]):
                sys.stdout.write('\rPlotting extinction histogram for species %i/%i'%(i+1,len(te_array[:,0])))
                plt.figure()
                species_te_array = te_array[:,1:][i]
                not_na_values = species_te_array[~np.isnan(list(species_te_array))]
                heights, bins = np.histogram(not_na_values,np.arange(0,delta_t+10,10))
                percent = heights/n_sim
                plt.bar(bins[:-1],percent,width=10, align="edge")
                plt.ylim(0,0.5)
                #survival_prob = 1-sum(percent)
                #if survival_prob >= 0.5:
                #    text_color = 'green'
                #else:
                #    text_color = 'red'
                ax = plt.gca()
                #plt.text(0.05, 0.7, 'survival probability: %.2f'%survival_prob,color=text_color, horizontalalignment='left',verticalalignment='baseline', transform=ax.transAxes)
                # annotate last bar
                if ax.patches[-1].get_height() > 0:
                    ax.text(ax.patches[-1].get_x()+3, np.round(ax.patches[-1].get_height()+0.001,4), '**', fontsize=12, color='black')
                plt.title('%s - Extinct in %i years: %i/%i'%(species,delta_t,sum(heights),n_sim))
                plt.xlabel('Years from present')
                plt.ylabel('Fraction of simulations')
                plt.tight_layout()
                pdf.savefig()  # saves the current figure into a pdf page
                plt.close()
    print('\n')
    #__________________________________________________________________________


# Copyright 2002 Gary Strangman.  All rights reserved
# Copyright 2002-2016 The SciPy Developers
#
# The original code from Gary Strangman was heavily adapted for
# use in SciPy by Travis Oliphant.  The original code came with the
# following disclaimer:
#
# This software is provided "as-is".  There are no expressed or implied
# warranties of any kind, including, but not limited to, the warranties
# of merchantability and fitness for a given application.  In no event
# shall Gary Strangman be liable for any direct, indirect, incidental,
# special, exemplary or consequential damages (including, but not limited
# to, loss of use, data or profits, or business interruption) however
# caused and on any theory of liability, whether in contract, strict
# liability or tort (including negligence or otherwise) arising in any way
# out of the use of this software, even if advised of the possibility of
# such damage.

"""
A collection of basic statistical functions for Python.

References
----------
.. [CRCProbStat2000] Zwillinger, D. and Kokoska, S. (2000). CRC Standard
   Probability and Statistics Tables and Formulae. Chapman & Hall: New
   York. 2000.

"""
import warnings
import math
from math import gcd
from collections import namedtuple

import numpy as np
from numpy import array, asarray, ma

from scipy.spatial.distance import cdist
from scipy.ndimage import measurements
from scipy._lib._util import (check_random_state, MapWrapper,
                              rng_integers, float_factorial)
import scipy.special as special
from scipy import linalg
from scipy.stats import distributions

# Functions/classes in other files should be added in `__init__.py`, not here
__all__ = ['find_repeats', 'gmean', 'hmean', 'mode', 'tmean', 'tvar',
           'tmin', 'tmax', 'tstd', 'tsem', 'moment', 'variation',
           'skew', 'kurtosis', 'describe', 'skewtest', 'kurtosistest',
           'normaltest', 'jarque_bera', 'itemfreq',
           'scoreatpercentile', 'percentileofscore',
           'cumfreq', 'relfreq', 'obrientransform',
           'sem', 'zmap', 'zscore', 'iqr', 'gstd', 'median_absolute_deviation',
           'median_abs_deviation',
           'sigmaclip', 'trimboth', 'trim1', 'trim_mean',
           'f_oneway', 'F_onewayConstantInputWarning',
           'F_onewayBadInputSizesWarning',
           'PearsonRConstantInputWarning', 'PearsonRNearConstantInputWarning',
           'pearsonr', 'fisher_exact',
           'SpearmanRConstantInputWarning', 'spearmanr', 'pointbiserialr',
           'kendalltau', 'weightedtau', 'multiscale_graphcorr',
           'linregress', 'siegelslopes', 'theilslopes', 'ttest_1samp',
           'ttest_ind', 'ttest_ind_from_stats', 'ttest_rel',
           'kstest', 'ks_1samp', 'ks_2samp',
           'chisquare', 'power_divergence',
           'tiecorrect', 'ranksums', 'kruskal', 'friedmanchisquare',
           'rankdata',
           'combine_pvalues', 'wasserstein_distance', 'energy_distance',
           'brunnermunzel', 'alexandergovern']

def _contains_nan(a, nan_policy='propagate'):
    policies = ['propagate', 'raise', 'omit']
    if nan_policy not in policies:
        raise ValueError("nan_policy must be one of {%s}" %
                         ', '.join("'%s'" % s for s in policies))
    try:
        # Calling np.sum to avoid creating a huge array into memory
        # e.g. np.isnan(a).any()
        with np.errstate(invalid='ignore'):
            contains_nan = np.isnan(np.sum(a))
    except TypeError:
        # This can happen when attempting to sum things which are not
        # numbers (e.g. as in the function `mode`). Try an alternative method:
        try:
            contains_nan = np.nan in set(a.ravel())
        except TypeError:
            # Don't know what to do. Fall back to omitting nan values and
            # issue a warning.
            contains_nan = False
            nan_policy = 'omit'
            warnings.warn("The input array could not be properly "
                          "checked for nan values. nan values "
                          "will be ignored.", RuntimeWarning)

    if contains_nan and nan_policy == 'raise':
        raise ValueError("The input contains nan values")

    return contains_nan, nan_policy

def _chk_asarray(a, axis):
    if axis is None:
        a = np.ravel(a)
        outaxis = 0
    else:
        a = np.asarray(a)
        outaxis = axis

    if a.ndim == 0:
        a = np.atleast_1d(a)

    return a, outaxis

def _chk2_asarray(a, b, axis):
    if axis is None:
        a = np.ravel(a)
        b = np.ravel(b)
        outaxis = 0
    else:
        a = np.asarray(a)
        b = np.asarray(b)
        outaxis = axis

    if a.ndim == 0:
        a = np.atleast_1d(a)
    if b.ndim == 0:
        b = np.atleast_1d(b)

    return a, b, outaxis

def _shape_with_dropped_axis(a, axis):
    """
    Given an array `a` and an integer `axis`, return the shape
    of `a` with the `axis` dimension removed.

    Examples
    --------
    >>> a = np.zeros((3, 5, 2))
    >>> _shape_with_dropped_axis(a, 1)
    (3, 2)

    """
    shp = list(a.shape)
    try:
        del shp[axis]
    except IndexError:
        raise np.AxisError(axis, a.ndim) from None
    return tuple(shp)

def _broadcast_shapes(shape1, shape2):
    """
    Given two shapes (i.e. tuples of integers), return the shape
    that would result from broadcasting two arrays with the given
    shapes.

    Examples
    --------
    >>> _broadcast_shapes((2, 1), (4, 1, 3))
    (4, 2, 3)
    """
    d = len(shape1) - len(shape2)
    if d <= 0:
        shp1 = (1,)*(-d) + shape1
        shp2 = shape2
    else:
        shp1 = shape1
        shp2 = (1,)*d + shape2
    shape = []
    for n1, n2 in zip(shp1, shp2):
        if n1 == 1:
            n = n2
        elif n2 == 1 or n1 == n2:
            n = n1
        else:
            raise ValueError(f'shapes {shape1} and {shape2} could not be '
                             'broadcast together')
        shape.append(n)
    return tuple(shape)

def _broadcast_shapes_with_dropped_axis(a, b, axis):
    """
    Given two arrays `a` and `b` and an integer `axis`, find the
    shape of the broadcast result after dropping `axis` from the
    shapes of `a` and `b`.

    Examples
    --------
    >>> a = np.zeros((5, 2, 1))
    >>> b = np.zeros((1, 9, 3))
    >>> _broadcast_shapes_with_dropped_axis(a, b, 1)
    (5, 3)
    """
    shp1 = _shape_with_dropped_axis(a, axis)
    shp2 = _shape_with_dropped_axis(b, axis)
    try:
        shp = _broadcast_shapes(shp1, shp2)
    except ValueError:
        raise ValueError(f'non-axis shapes {shp1} and {shp2} could not be '
                         'broadcast together') from None
    return shp

# Map from names to lambda_ values used in power_divergence().
_power_div_lambda_names = {
    "pearson": 1,
    "log-likelihood": 0,
    "freeman-tukey": -0.5,
    "mod-log-likelihood": -1,
    "neyman": -2,
    "cressie-read": 2/3,
}

def _count(a, axis=None):
    """Count the number of non-masked elements of an array.

    This function behaves like `np.ma.count`, but is much faster
    for ndarrays.
    """
    if hasattr(a, 'count'):
        num = a.count(axis=axis)
        if isinstance(num, np.ndarray) and num.ndim == 0:
            # In some cases, the `count` method returns a scalar array (e.g.
            # np.array(3)), but we want a plain integer.
            num = int(num)
    else:
        if axis is None:
            num = a.size
        else:
            num = a.shape[axis]
    return num

def _m_broadcast_to(a, shape):
    if np.ma.isMaskedArray(a):
        return np.ma.masked_array(np.broadcast_to(a, shape),
                                  mask=np.broadcast_to(a.mask, shape))
    return np.broadcast_to(a, shape, subok=True)

Power_divergenceResult = namedtuple('Power_divergenceResult',
                                    ('statistic', 'pvalue'))

def power_divergence(f_obs, f_exp=None, ddof=0, axis=0, lambda_=None):
    """Cressie-Read power divergence statistic and goodness of fit test.

    This function tests the null hypothesis that the categorical data
    has the given frequencies, using the Cressie-Read power divergence
    statistic.

    Parameters
    ----------
    f_obs : array_like
        Observed frequencies in each category.
    f_exp : array_like, optional
        Expected frequencies in each category.  By default the categories are
        assumed to be equally likely.
    ddof : int, optional
        "Delta degrees of freedom": adjustment to the degrees of freedom
        for the p-value.  The p-value is computed using a chi-squared
        distribution with ``k - 1 - ddof`` degrees of freedom, where `k`
        is the number of observed frequencies.  The default value of `ddof`
        is 0.
    axis : int or None, optional
        The axis of the broadcast result of `f_obs` and `f_exp` along which to
        apply the test.  If axis is None, all values in `f_obs` are treated
        as a single data set.  Default is 0.
    lambda_ : float or str, optional
        The power in the Cressie-Read power divergence statistic.  The default
        is 1.  For convenience, `lambda_` may be assigned one of the following
        strings, in which case the corresponding numerical value is used::

            String              Value   Description
            "pearson"             1     Pearson's chi-squared statistic.
                                        In this case, the function is
                                        equivalent to `stats.chisquare`.
            "log-likelihood"      0     Log-likelihood ratio. Also known as
                                        the G-test [3]_.
            "freeman-tukey"      -1/2   Freeman-Tukey statistic.
            "mod-log-likelihood" -1     Modified log-likelihood ratio.
            "neyman"             -2     Neyman's statistic.
            "cressie-read"        2/3   The power recommended in [5]_.

    Returns
    -------
    statistic : float or ndarray
        The Cressie-Read power divergence test statistic.  The value is
        a float if `axis` is None or if` `f_obs` and `f_exp` are 1-D.
    pvalue : float or ndarray
        The p-value of the test.  The value is a float if `ddof` and the
        return value `stat` are scalars.

    See Also
    --------
    chisquare

    Notes
    -----
    This test is invalid when the observed or expected frequencies in each
    category are too small.  A typical rule is that all of the observed
    and expected frequencies should be at least 5.

    Also, the sum of the observed and expected frequencies must be the same
    for the test to be valid; `power_divergence` raises an error if the sums
    do not agree within a relative tolerance of ``1e-8``.

    When `lambda_` is less than zero, the formula for the statistic involves
    dividing by `f_obs`, so a warning or error may be generated if any value
    in `f_obs` is 0.

    Similarly, a warning or error may be generated if any value in `f_exp` is
    zero when `lambda_` >= 0.

    The default degrees of freedom, k-1, are for the case when no parameters
    of the distribution are estimated. If p parameters are estimated by
    efficient maximum likelihood then the correct degrees of freedom are
    k-1-p. If the parameters are estimated in a different way, then the
    dof can be between k-1-p and k-1. However, it is also possible that
    the asymptotic distribution is not a chisquare, in which case this
    test is not appropriate.

    This function handles masked arrays.  If an element of `f_obs` or `f_exp`
    is masked, then data at that position is ignored, and does not count
    towards the size of the data set.

    .. versionadded:: 0.13.0

    References
    ----------
    .. [1] Lowry, Richard.  "Concepts and Applications of Inferential
           Statistics". Chapter 8.
           https://web.archive.org/web/20171015035606/http://faculty.vassar.edu/lowry/ch8pt1.html
    .. [2] "Chi-squared test", https://en.wikipedia.org/wiki/Chi-squared_test
    .. [3] "G-test", https://en.wikipedia.org/wiki/G-test
    .. [4] Sokal, R. R. and Rohlf, F. J. "Biometry: the principles and
           practice of statistics in biological research", New York: Freeman
           (1981)
    .. [5] Cressie, N. and Read, T. R. C., "Multinomial Goodness-of-Fit
           Tests", J. Royal Stat. Soc. Series B, Vol. 46, No. 3 (1984),
           pp. 440-464.

    Examples
    --------
    (See `chisquare` for more examples.)

    When just `f_obs` is given, it is assumed that the expected frequencies
    are uniform and given by the mean of the observed frequencies.  Here we
    perform a G-test (i.e. use the log-likelihood ratio statistic):

    >>> from scipy.stats import power_divergence
    >>> power_divergence([16, 18, 16, 14, 12, 12], lambda_='log-likelihood')
    (2.006573162632538, 0.84823476779463769)

    The expected frequencies can be given with the `f_exp` argument:

    >>> power_divergence([16, 18, 16, 14, 12, 12],
    ...                  f_exp=[16, 16, 16, 16, 16, 8],
    ...                  lambda_='log-likelihood')
    (3.3281031458963746, 0.6495419288047497)

    When `f_obs` is 2-D, by default the test is applied to each column.

    >>> obs = np.array([[16, 18, 16, 14, 12, 12], [32, 24, 16, 28, 20, 24]]).T
    >>> obs.shape
    (6, 2)
    >>> power_divergence(obs, lambda_="log-likelihood")
    (array([ 2.00657316,  6.77634498]), array([ 0.84823477,  0.23781225]))

    By setting ``axis=None``, the test is applied to all data in the array,
    which is equivalent to applying the test to the flattened array.

    >>> power_divergence(obs, axis=None)
    (23.31034482758621, 0.015975692534127565)
    >>> power_divergence(obs.ravel())
    (23.31034482758621, 0.015975692534127565)

    `ddof` is the change to make to the default degrees of freedom.

    >>> power_divergence([16, 18, 16, 14, 12, 12], ddof=1)
    (2.0, 0.73575888234288467)

    The calculation of the p-values is done by broadcasting the
    test statistic with `ddof`.

    >>> power_divergence([16, 18, 16, 14, 12, 12], ddof=[0,1,2])
    (2.0, array([ 0.84914504,  0.73575888,  0.5724067 ]))

    `f_obs` and `f_exp` are also broadcast.  In the following, `f_obs` has
    shape (6,) and `f_exp` has shape (2, 6), so the result of broadcasting
    `f_obs` and `f_exp` has shape (2, 6).  To compute the desired chi-squared
    statistics, we must use ``axis=1``:

    >>> power_divergence([16, 18, 16, 14, 12, 12],
    ...                  f_exp=[[16, 16, 16, 16, 16, 8],
    ...                         [8, 20, 20, 16, 12, 12]],
    ...                  axis=1)
    (array([ 3.5 ,  9.25]), array([ 0.62338763,  0.09949846]))

    """
    # Convert the input argument `lambda_` to a numerical value.
    if isinstance(lambda_, str):
        if lambda_ not in _power_div_lambda_names:
            names = repr(list(_power_div_lambda_names.keys()))[1:-1]
            raise ValueError("invalid string for lambda_: {0!r}. "
                             "Valid strings are {1}".format(lambda_, names))
        lambda_ = _power_div_lambda_names[lambda_]
    elif lambda_ is None:
        lambda_ = 1

    f_obs = np.asanyarray(f_obs)
    f_obs_float = f_obs.astype(np.float64)

    if f_exp is not None:
        f_exp = np.asanyarray(f_exp)
        bshape = _broadcast_shapes(f_obs_float.shape, f_exp.shape)
        f_obs_float = _m_broadcast_to(f_obs_float, bshape)
        f_exp = _m_broadcast_to(f_exp, bshape)
        rtol = 1e-2  # to pass existing tests #EDIT from 1e-8 to 1e-2
        with np.errstate(invalid='ignore'):
            f_obs_sum = f_obs_float.sum(axis=axis)
            f_exp_sum = f_exp.sum(axis=axis)
            relative_diff = (np.abs(f_obs_sum - f_exp_sum) /
                             np.minimum(f_obs_sum, f_exp_sum))
            diff_gt_tol = (relative_diff > rtol).any()
        if diff_gt_tol:
            msg = (f"For each axis slice, the sum of the observed "
                   f"frequencies must agree with the sum of the "
                   f"expected frequencies to a relative tolerance "
                   f"of {rtol}, but the percent differences are:\n"
                   f"{relative_diff}")
            raise ValueError(msg)

    else:
        # Ignore 'invalid' errors so the edge case of a data set with length 0
        # is handled without spurious warnings.
        with np.errstate(invalid='ignore'):
            f_exp = f_obs.mean(axis=axis, keepdims=True)

    # `terms` is the array of terms that are summed along `axis` to create
    # the test statistic.  We use some specialized code for a few special
    # cases of lambda_.
    if lambda_ == 1:
        # Pearson's chi-squared statistic
        terms = (f_obs_float - f_exp)**2 / f_exp
    elif lambda_ == 0:
        # Log-likelihood ratio (i.e. G-test)
        terms = 2.0 * special.xlogy(f_obs, f_obs / f_exp)
    elif lambda_ == -1:
        # Modified log-likelihood ratio
        terms = 2.0 * special.xlogy(f_exp, f_exp / f_obs)
    else:
        # General Cressie-Read power divergence.
        terms = f_obs * ((f_obs / f_exp)**lambda_ - 1)
        terms /= 0.5 * lambda_ * (lambda_ + 1)

    stat = terms.sum(axis=axis)

    num_obs = _count(terms, axis=axis)
    ddof = asarray(ddof)
    p = distributions.chi2.sf(stat, num_obs - 1 - ddof)

    return Power_divergenceResult(stat, p)

def chisquare(f_obs, f_exp=None, ddof=0, axis=0):
    """Calculate a one-way chi-square test.

    The chi-square test tests the null hypothesis that the categorical data
    has the given frequencies.

    Parameters
    ----------
    f_obs : array_like
        Observed frequencies in each category.
    f_exp : array_like, optional
        Expected frequencies in each category.  By default the categories are
        assumed to be equally likely.
    ddof : int, optional
        "Delta degrees of freedom": adjustment to the degrees of freedom
        for the p-value.  The p-value is computed using a chi-squared
        distribution with ``k - 1 - ddof`` degrees of freedom, where `k`
        is the number of observed frequencies.  The default value of `ddof`
        is 0.
    axis : int or None, optional
        The axis of the broadcast result of `f_obs` and `f_exp` along which to
        apply the test.  If axis is None, all values in `f_obs` are treated
        as a single data set.  Default is 0.

    Returns
    -------
    chisq : float or ndarray
        The chi-squared test statistic.  The value is a float if `axis` is
        None or `f_obs` and `f_exp` are 1-D.
    p : float or ndarray
        The p-value of the test.  The value is a float if `ddof` and the
        return value `chisq` are scalars.

    See Also
    --------
    scipy.stats.power_divergence
    scipy.stats.fisher_exact : Fisher exact test on a 2x2 contingency table.
    scipy.stats.barnard_exact : An unconditional exact test. An alternative
        to chi-squared test for small sample sizes.

    Notes
    -----
    This test is invalid when the observed or expected frequencies in each
    category are too small.  A typical rule is that all of the observed
    and expected frequencies should be at least 5. According to [3]_, the
    total number of samples is recommended to be greater than 13,
    otherwise exact tests (such as Barnard's Exact test) should be used
    because they do not overreject.

    Also, the sum of the observed and expected frequencies must be the same
    for the test to be valid; `chisquare` raises an error if the sums do not
    agree within a relative tolerance of ``1e-8``.

    The default degrees of freedom, k-1, are for the case when no parameters
    of the distribution are estimated. If p parameters are estimated by
    efficient maximum likelihood then the correct degrees of freedom are
    k-1-p. If the parameters are estimated in a different way, then the
    dof can be between k-1-p and k-1. However, it is also possible that
    the asymptotic distribution is not chi-square, in which case this test
    is not appropriate.

    References
    ----------
    .. [1] Lowry, Richard.  "Concepts and Applications of Inferential
           Statistics". Chapter 8.
           https://web.archive.org/web/20171022032306/http://vassarstats.net:80/textbook/ch8pt1.html
    .. [2] "Chi-squared test", https://en.wikipedia.org/wiki/Chi-squared_test
    .. [3] Pearson, Karl. "On the criterion that a given system of deviations from the probable
           in the case of a correlated system of variables is such that it can be reasonably
           supposed to have arisen from random sampling", Philosophical Magazine. Series 5. 50
           (1900), pp. 157-175.

    Examples
    --------
    When just `f_obs` is given, it is assumed that the expected frequencies
    are uniform and given by the mean of the observed frequencies.

    >>> from scipy.stats import chisquare
    >>> chisquare([16, 18, 16, 14, 12, 12])
    (2.0, 0.84914503608460956)

    With `f_exp` the expected frequencies can be given.

    >>> chisquare([16, 18, 16, 14, 12, 12], f_exp=[16, 16, 16, 16, 16, 8])
    (3.5, 0.62338762774958223)

    When `f_obs` is 2-D, by default the test is applied to each column.

    >>> obs = np.array([[16, 18, 16, 14, 12, 12], [32, 24, 16, 28, 20, 24]]).T
    >>> obs.shape
    (6, 2)
    >>> chisquare(obs)
    (array([ 2.        ,  6.66666667]), array([ 0.84914504,  0.24663415]))

    By setting ``axis=None``, the test is applied to all data in the array,
    which is equivalent to applying the test to the flattened array.

    >>> chisquare(obs, axis=None)
    (23.31034482758621, 0.015975692534127565)
    >>> chisquare(obs.ravel())
    (23.31034482758621, 0.015975692534127565)

    `ddof` is the change to make to the default degrees of freedom.

    >>> chisquare([16, 18, 16, 14, 12, 12], ddof=1)
    (2.0, 0.73575888234288467)

    The calculation of the p-values is done by broadcasting the
    chi-squared statistic with `ddof`.

    >>> chisquare([16, 18, 16, 14, 12, 12], ddof=[0,1,2])
    (2.0, array([ 0.84914504,  0.73575888,  0.5724067 ]))

    `f_obs` and `f_exp` are also broadcast.  In the following, `f_obs` has
    shape (6,) and `f_exp` has shape (2, 6), so the result of broadcasting
    `f_obs` and `f_exp` has shape (2, 6).  To compute the desired chi-squared
    statistics, we use ``axis=1``:

    >>> chisquare([16, 18, 16, 14, 12, 12],
    ...           f_exp=[[16, 16, 16, 16, 16, 8], [8, 20, 20, 16, 12, 12]],
    ...           axis=1)
    (array([ 3.5 ,  9.25]), array([ 0.62338763,  0.09949846]))

    """
    return power_divergence(f_obs, f_exp=f_exp, ddof=ddof, axis=axis,
                            lambda_="pearson")

def _sum_of_squares(a, axis=0):
    """Square each element of the input array, and return the sum(s) of that.

    Parameters
    ----------
    a : array_like
        Input array.
    axis : int or None, optional
        Axis along which to calculate. Default is 0. If None, compute over
        the whole array `a`.

    Returns
    -------
    sum_of_squares : ndarray
        The sum along the given axis for (a**2).

    See Also
    --------
    _square_of_sums : The square(s) of the sum(s) (the opposite of
        `_sum_of_squares`).

    """
    a, axis = _chk_asarray(a, axis)
    return np.sum(a*a, axis)

def _square_of_sums(a, axis=0):
    """Sum elements of the input array, and return the square(s) of that sum.

    Parameters
    ----------
    a : array_like
        Input array.
    axis : int or None, optional
        Axis along which to calculate. Default is 0. If None, compute over
        the whole array `a`.

    Returns
    -------
    square_of_sums : float or ndarray
        The square of the sum over `axis`.

    See Also
    --------
    _sum_of_squares : The sum of squares (the opposite of `square_of_sums`).

    """
    a, axis = _chk_asarray(a, axis)
    s = np.sum(a, axis)
    if not np.isscalar(s):
        return s.astype(float) * s
    else:
        return float(s) * s

import numpy,from numpy,import scipy,from scipy


#!/usr/bin/env python
#This script plots drag around an inline oscillating cylinder for re 200 kc 10 against dutsch et als work at cycle 14
import argparse
import os
import os.path
import sys
import csv
import matplotlib
from matplotlib import pyplot as plt
import numpy

cuibmFolder = os.path.expandvars("/scratch/src/cuIBM")

validationData = '/osc_Re200_KC10_Dutsch.txt'

execPath       = cuibmFolder + '/bin/cuIBM'
caseFolder     = cuibmFolder + '/validation/osc/static'
validationData = cuibmFolder + '/validation-data' + validationData

print "\n"+"-"*100
print "Plotting validation for flow around inline oscillating cylinder with Re200 and KC10\n"
print "-"*100+"\n"

experiment = numpy.genfromtxt(validationData,delimiter='\t')
external = numpy.genfromtxt(caseFolder+'/externalkc10/forces',delimiter='\t')
embedded = numpy.genfromtxt(caseFolder+'/embeddedkc10/forces',delimiter='\t')

#external
plt.plot([i-13 for i in zip(*external)[0]],[i*5 for i in zip(*external)[1]],'-',color='blue',linewidth=2,label='External')
plt.plot(zip(*experiment)[0],zip(*experiment)[1],'o', color = 'red', markersize = 8, label = 'Dutsch et al')
plt.title('Drag for flow around inline oscillating cylinder Re 200, KC 10')
plt.legend(loc='lower right',numpoints=1, fancybox=True)
plt.xlabel('t/T')
plt.ylabel('Fd')
plt.ylim([-6,6])
plt.xlim([0,1])
plt.savefig('%s/External_static_kc10.pdf' % (caseFolder))
plt.clf()

#emb
plt.plot([i-13 for i in zip(*embedded)[0]],[i*5 for i in zip(*embedded)[1]],'-',color='blue',linewidth=2,label='Embedded')
plt.plot(zip(*experiment)[0],zip(*experiment)[1],'o', color = 'red', markersize = 8, label = 'Dutsch et al')
plt.title('Drag for flow around inline oscillating cylinder Re 200, KC 10')
plt.legend(loc='lower right',numpoints=1, fancybox=True)
plt.xlabel('t/T')
plt.ylabel('Fd')
plt.ylim([-6,6])
plt.xlim([0,1])
plt.savefig('%s/Embedded_static_kc10.pdf' % (caseFolder))
plt.clf()

print '\nDone plotting!\n Files saved to %s' % caseFolder


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

"""RGB color space.
"""

from __future__ import annotations

from typing import cast
from typing import Union

import cv2
import numpy as np
import torch
from multipledispatch import dispatch
from torch import nn
from torch import Tensor

from onevision.cv.utils import batch_image_processing
from onevision.cv.utils import channel_last_processing
from onevision.factory import TRANSFORMS
from onevision.type import TensorOrArray

__all__ = [
    "bgr_to_rgb",
    "bgr_to_rgba",
    "linear_rgb_to_rgb",
    "rgb_to_bgr",
    "rgb_to_linear_rgb",
    "rgb_to_rgba",
    "rgba_to_bgr",
    "rgba_to_rgb",
    "BgrToRgb",
    "BgrToRgba",
    "LinearRgbToRgb",
    "RgbaToBgr",
    "RgbaToRgb",
    "RgbToBgr",
    "RgbToLinearRgb",
    "RgbToRgba"
]

# MARK: - Functional

@dispatch(Tensor)
def bgr_to_rgb(image: Tensor) -> Tensor:
    """Convert a BGR image to RGB.

    Args:
        image (Tensor[B, 3, H, W]):
            BGR Image to be converted to BGR.

    Returns:
        rgb (Tensor[B, 3, H, W]):
            RGB version of the image.
    """
    # Flip image channels
    rgb = image.flip(-3)
    return rgb

@batch_image_processing
@channel_last_processing
@dispatch(np.ndarray)
def bgr_to_rgb(image: np.ndarray) -> np.ndarray:
    """Convert a BGR image to RGB.

    Args:
        image (np.ndarray[B, 3, H, W]):
            BGR Image to be converted to BGR.

    Returns:
        rgb (np.ndarray[B, 3, H, W]):
            RGB version of the image.
    """
    return cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

@dispatch(Tensor, (float, Tensor))
def bgr_to_rgba(image: Tensor, alpha_val: Union[float, Tensor]) -> Tensor:
    """Convert an image from BGR to RGBA.

    Args:
        image (Tensor[B, 3, H, W]):
            BGR Image to be converted to RGBA.
        alpha_val (float, Tensor[B, 1, H, W]):
            A float number or tensor for the alpha value.

    Returns:
        rgba (Tensor[B, 4, H, W]):
            RGBA version of the image.

    Notes:
        Current functionality is NOT supported by Torchscript.
    """
    if not isinstance(alpha_val, (float, Tensor)):
        raise TypeError(f"`alpha_val` must be a `float` or `Tensor`. "
                        f"But got: {type(alpha_val)}.")

    # Convert first to RGB, then add alpha channel
    rgb  = bgr_to_rgb(image)
    rgba = rgb_to_rgba(rgb, alpha_val)

    return rgba

@batch_image_processing
@channel_last_processing
@dispatch(np.ndarray, (float, np.ndarray))
def bgr_to_rgba(image: np.ndarray, alpha_val: Union[float, np.ndarray]) -> np.ndarray:
    """Convert an image from BGR to RGBA.

    Args:
        image (np.ndarray[B, 3, H, W]):
            BGR Image to be converted to RGBA.
        alpha_val (float, np.ndarray[B, 1, H, W]):
            A float number or tensor for the alpha value.

    Returns:
        rgba (np.ndarray[B, 4, H, W]):
            RGBA version of the image.

    Notes:
        Current functionality is NOT supported by Torchscript.
    """
    return cv2.cvtColor(image, cv2.COLOR_BGR2RGBA)

@dispatch(Tensor)
def linear_rgb_to_rgb(image: Tensor) -> Tensor:
    """Convert a linear RGB image to sRGB. Used in colorspace conversions.

    Args:
        image (Tensor[B, 3, H, W]):
            Linear RGB Image to be converted to sRGB.

    Returns:
        rgb (Tensor[B, 3, H, W]):
            sRGB version of the image.
    """
    threshold = 0.0031308
    rgb       = torch.where(
        image > threshold,
        1.055 * torch.pow(image.clamp(min=threshold), 1 / 2.4) - 0.055,
        12.92 * image
    )
    return rgb

def rgb_to_bgr(image: TensorOrArray) -> TensorOrArray:
    """Convert an RGB image to BGR.

    Args:
        image (TensorOrArray[B, 3, H, W]):
            RGB Image to be converted to BGR.

    Returns:
        bgr (TensorOrArray[B, 3, H, W]):
            BGR version of the image.
    """
    return bgr_to_rgb(image)

@dispatch(Tensor)
def rgb_to_linear_rgb(image: Tensor) -> Tensor:
    """Convert an sRGB image to linear RGB. Used in colorspace conversions.

    Args:
        image (Tensor[B, 3, H, W]):
            sRGB Image to be converted to linear RGB.

    Returns:
        linear_rgb (Tensor[B, 3, H, W]):
            linear RGB version of the image.
    """
    lin_rgb = torch.where(
        image > 0.04045,
        torch.pow(((image + 0.055) / 1.055), 2.4),
        image / 12.92
    )
    return lin_rgb

@dispatch(Tensor, (float, Tensor))
def rgb_to_rgba(image: Tensor, alpha_val: Union[float, Tensor]) -> Tensor:
    """Convert an image from RGB to RGBA.

    Args:
        image (Tensor[B, 3, H, W]):
            RGB Image to be converted to RGBA.
        alpha_val (float, Tensor[B, 1, H, W]):
            A float number or tensor for the alpha value.

    Returns:
        rgba (Tensor[B, 4, H, W]):
            RGBA version of the image

    Notes:
        Current functionality is NOT supported by Torchscript.
    """
    if not isinstance(alpha_val, (float, Tensor)):
        raise TypeError(f"`alpha_val` must be `float` or `Tensor`. "
                        f"But got: {type(alpha_val)}.")

    # Add one channel
    r, g, b = torch.chunk(image, image.shape[-3], dim=-3)
    a       = cast(Tensor, alpha_val)

    if isinstance(alpha_val, float):
        a = torch.full_like(r, fill_value=float(alpha_val))
    rgba = torch.cat([r, g, b, a], dim=-3)

    return rgba

@batch_image_processing
@channel_last_processing
@dispatch(np.ndarray, (float, np.ndarray))
def rgb_to_rgba(image: np.ndarray, alpha_val: Union[float, np.ndarray]) -> np.ndarray:
    """Convert an image from RGB to RGBA.

    Args:
        image (np.ndarray[B, 3, H, W]):
            RGB Image to be converted to RGBA.
        alpha_val (float, np.ndarray[B, 1, H, W]):
            A float number or tensor for the alpha value.

    Returns:
        rgba (np.ndarray[B, 4, H, W]):
            RGBA version of the image

    Notes:
        Current functionality is NOT supported by Torchscript.
    """
    return cv2.cvtColor(image, cv2.COLOR_RGB2RGBA)

@dispatch(Tensor)
def rgba_to_bgr(image: Tensor) -> Tensor:
    """Convert an image from RGBA to BGR.

    Args:
        image (Tensor[B, 4, H, W]):
            RGBA Image to be converted to BGR.

    Returns:
        rgb (Tensor[B, 3, H, W]):
            RGB version of the image.
    """
    if not isinstance(image, Tensor):
        raise TypeError(f"`image` must be a `Tensor`. But got: {type(image)}")
    if image.ndim < 3 or image.shape[-3] != 4:
        raise ValueError(f"`image` must have a shape of [*, 4, H, W]. "
                         f"But got: {image.shape}")

    # Convert to RGB first, then to BGR
    rgb = rgba_to_rgb(image)
    bgr = rgb_to_bgr(rgb)

    return bgr

@batch_image_processing
@channel_last_processing
@dispatch(np.ndarray)
def rgba_to_bgr(image: np.ndarray) -> np.ndarray:
    """Convert an image from RGBA to BGR.

    Args:
        image (np.ndarray[B, 4, H, W]):
            RGBA Image to be converted to BGR.

    Returns:
        bgr (np.ndarray[B, 3, H, W]):
            BGR version of the image.
    """
    return cv2.cvtColor(image, cv2.COLOR_RGBA2BGR)

@dispatch(Tensor)
def rgba_to_rgb(image: Tensor) -> Tensor:
    """Convert an image from RGBA to RGB.

    Args:
        image (Tensor[B, 4, H, W]):
            RGBA Image to be converted to RGB.

    Returns:
        rgb (Tensor[B, 3, H, W]):
            RGB version of the image.
    """
    # Unpack channels
    r, g, b, a = torch.chunk(image, image.shape[-3], dim=-3)

    # Compute new channels
    a_one = torch.tensor(1.0) - a
    r_new = a_one * r + a * r
    g_new = a_one * g + a * g
    b_new = a_one * b + a * b
    rgb   = torch.cat([r_new, g_new, b_new], dim=-3)

    return rgb

@batch_image_processing
@channel_last_processing
@dispatch(np.ndarray)
def rgba_to_rgb(image: np.ndarray) -> np.ndarray:
    """Convert an image from RGBA to RGB.

    Args:
        image (np.ndarray[B, 4, H, W]):
            RGBA Image to be converted to RGB.

    Returns:
        rgb (np.ndarray[B, 3, H, W]):
            RGB version of the image.
    """
    return cv2.cvtColor(image, cv2.COLOR_RGBA2RGB)

# MARK: - Modules

@TRANSFORMS.register(name="bgr_to_rgb")
class BgrToRgb(nn.Module):
    """Convert image from BGR to RGB. Image data is assumed to be in the
    range of [0.0, 1.0].
    """

    # MARK: Forward Pass

    def forward(self, image: TensorOrArray) -> TensorOrArray:
        return bgr_to_rgb(image)

@TRANSFORMS.register(name="bgr_to_rgba")
class BgrToRgba(nn.Module):
    """Convert an image from BGR to RGBA. Add an alpha channel to existing RGB
    image.

    Args:
        alpha_val (float, TensorOrArray[B, 1, H, W]):
            A float number or tensor for the alpha value.

    Notes:
        Current functionality is NOT supported by Torchscript.
    """

    # MARK: Magic Functions

    def __init__(self, alpha_val: Union[float, TensorOrArray]):
        super().__init__()
        self.alpha_val = alpha_val

    # MARK: Forward Pass

    def forward(self, image: TensorOrArray) -> TensorOrArray:
        return rgb_to_rgba(image, self.alpha_val)

@TRANSFORMS.register(name="linear_rgb_to_rgb")
class LinearRgbToRgb(nn.Module):
    """Convert a linear RGB image to sRGB. Applies gamma correction to linear
    RGB values, at the end of colorspace conversions, to get sRGB.

    References:
        [1] https://stackoverflow.com/questions/35952564/convert-rgb-to-srgb
        [2] https://www.cambridgeincolour.com/tutorials/gamma-correction.htm
        [3] https://en.wikipedia.org/wiki/SRGB
    """

    # MARK: Forward Pass

    def forward(self, image: TensorOrArray) -> TensorOrArray:
        return linear_rgb_to_rgb(image)

@TRANSFORMS.register(name="rgb_to_bgr")
class RgbToBgr(nn.Module):
    """Convert an image from RGB to BGR. Image data is assumed to be in the
    range of [0.0, 1.0].
    """

    # MARK: Forward Pass

    def forward(self, image: TensorOrArray) -> TensorOrArray:
        return rgb_to_bgr(image)

@TRANSFORMS.register(name="rgb_to_linear_rgb")
class RgbToLinearRgb(nn.Module):
    """Convert an image from sRGB to linear RGB. Reverses the gamma correction
    of sRGB to get linear RGB values for colorspace conversions. Image data
    is assumed to be in the range of [0.0, 1.0].

    References:
        [1] https://stackoverflow.com/questions/35952564/convert-rgb-to-srgb
        [2] https://www.cambridgeincolour.com/tutorials/gamma-correction.htm
        [3] https://en.wikipedia.org/wiki/SRGB
    """

    # MARK: Forward Pass

    def forward(self, image: TensorOrArray) -> TensorOrArray:
        return rgb_to_linear_rgb(image)

@TRANSFORMS.register(name="rgb_to_rgba")
class RgbToRgba(nn.Module):
    """Convert an image from RGB to RGBA. Add an alpha channel to existing RGB
    image.

    Args:
        alpha_val (float, TensorOrArray[B, 1, H, W]):
            A float number or tensor for the alpha value.

    Notes:
        Current functionality is NOT supported by Torchscript.
    """

    # MARK: Magic Functions

    def __init__(self, alpha_val: Union[float, TensorOrArray]) -> None:
        super().__init__()
        self.alpha_val = alpha_val

    # MARK: Forward Pass

    def forward(self, image: TensorOrArray) -> TensorOrArray:
        return rgb_to_rgba(image, self.alpha_val)

@TRANSFORMS.register(name="rgba_to_bgr")
class RgbaToBgr(nn.Module):
    """Convert an image from RGBA to BGR. Remove an alpha channel from BGR
    image.
    """

    # MARK: Forward Pass

    def forward(self, image: TensorOrArray) -> TensorOrArray:
        return rgba_to_bgr(image)

@TRANSFORMS.register(name="rgba_to_rgb")
class RgbaToRgb(nn.Module):
    """Convert an image from RGBA to RGB. Remove an alpha channel from RGB
    image.
    """

    # MARK: Forward Pass

    def forward(self, image: TensorOrArray) -> TensorOrArray:
        return rgba_to_rgb(image)


import math
from typing import Optional, Union, Tuple, List

import numpy as np

import torch
import torch.nn.functional as F
from torch import Tensor
from torch_geometric.utils.num_nodes import maybe_num_nodes
from torch_scatter import scatter, segment_csr, gather_csr
from torch_scatter.utils import broadcast

import tsl

__all__ = [
    'expand_then_cat',
    'gated_tanh',
    'reverse_tensor',
    'sparse_softmax',
    'sparse_multi_head_attention'
]

def expand_then_cat(tensors: Union[Tuple[Tensor, ...], List[Tensor]],
                    dim=-1) -> Tensor:
    r"""
    Match the dimensions of tensors in the input list and then concatenate.

    Args:
        tensors: Tensors to concatenate.
        dim (int): Dimension along which to concatenate.
    """
    shapes = [t.shape for t in tensors]
    expand_dims = list(np.max(shapes, 0))
    expand_dims[dim] = -1
    tensors = [t.expand(*expand_dims) for t in tensors]
    return torch.cat(tensors, dim=dim)

@torch.jit.script
def gated_tanh(input: Tensor, dim: int = -1) -> Tensor:
    r"""The gated tanh unite. Computes:

    .. math ::
        \text{GatedTanH}(a, b) = \text{TanH}(a) \otimes \sigma(b)

    where `input` is split in half along `dim` to form `a` and `b`, :math:`\text{TanH}` is the hyperbolic tangent
    function, :math:`\sigma` is the sigmoid function and :math:`\otimes` is the element-wise product between matrices.

    Args:
        input (Tensor): Input tensor.
        dim (int, optional): Dimension on which to split the input.
                             (default: -1)
    """

    out, gate = torch.tensor_split(input, 2, dim=dim)
    return torch.tanh(out) * torch.sigmoid(gate)

@torch.jit.script
def reverse_tensor(tensor: Tensor, dim: int) -> Tensor:
    """Reverse tensor along specific dimension.

    Args:
        tensor (Tensor): Input tensor.
        dim (int): Dimension along which to reverse sequence.
    """
    indices = torch.arange(tensor.size(dim) - 1, -1, -1, device=tensor.device)
    return tensor.index_select(dim, indices)

@torch.jit.script
def sparse_softmax(src: Tensor, index: Optional[Tensor] = None,
                   ptr: Optional[Tensor] = None,
                   num_nodes: Optional[int] = None,
                   dim: int = -2) -> Tensor:
    r"""Extension of ~torch_geometric.softmax with index broadcasting to compute
    a sparsely evaluated softmax over multiple broadcast dimensions.

    Given a value tensor :attr:`src`, this function first groups the values
    along the first dimension based on the indices specified in :attr:`index`,
    and then proceeds to compute the softmax individually for each group.

    Args:
        src (Tensor): The source tensor.
        index (Tensor, optional): The indices of elements for applying the softmax.
        ptr (LongTensor, optional): If given, computes the softmax based on
            sorted inputs in CSR representation. (default: :obj:`None`)
        num_nodes (int, optional): The number of nodes, *i.e.*
            :obj:`max_val + 1` of :attr:`index`. (default: :obj:`None`)
        dim (int, optional): The dimension in which to normalize, i.e., the edge
            dimension. (default: :obj:`-2`)
    """
    if ptr is not None:
        dim = dim + src.dim() if dim < 0 else dim
        size = ([1] * dim) + [-1]
        ptr = ptr.view(size)
        src_max = gather_csr(segment_csr(src, ptr, reduce='max'), ptr)
        out = (src - src_max).exp()
        out_sum = gather_csr(segment_csr(out, ptr, reduce='sum'), ptr)
    elif index is not None:
        N = maybe_num_nodes(index, num_nodes)
        expanded_index = broadcast(index, src, dim)
        src_max = scatter(src, expanded_index, dim, dim_size=N, reduce='max')
        src_max = src_max.index_select(dim, index)
        out = (src - src_max).exp()
        out_sum = scatter(out, expanded_index, dim, dim_size=N, reduce='sum')
        out_sum = out_sum.index_select(dim, index)
    else:
        raise NotImplementedError

    return out / (out_sum + tsl.epsilon)

@torch.jit.script
def sparse_multi_head_attention(q: Tensor, k: Tensor, v: Tensor, index: Tensor,
                                dim_size: Optional[int] = None,
                                dropout_p: float = 0.0):
    r"""Computes multi-head, scaled, dot product attention on query, key and
    value tensors, applying dropout if a probability greater than 0.0 is
    specified. Index specifies for each query in q the belonging sequence in the
    original batched, dense tensor.
    Returns a tensor pair containing attended values and attention weights.

    Args:
        q (Tensor): Query tensor. See Shape section for shape details.
        k (Tensor): Key tensor. See Shape section for shape details.
        v (Tensor): Value tensor. See Shape section for shape details.
        index (Tensor): Tensor containing mask values to be added to calculated
            attention. May be 2D or 3D; see Shape section for details.
        dim_size (int, optional): The batched target length sequence, i.e.
            :obj:`max_val + 1` of :attr:`index`. (default: :obj:`None`)
        dropout_p: dropout probability. If greater than 0.0, dropout is applied.

    Shape:
        - q: :math:`(S, H, E)` where S is sparsed dimension, H is the number of
            heads, and E is embedding dimension.
        - k: :math:`(S, H, E)` where S is sparsed dimension, H is the number of
            heads, and E is embedding dimension.
        - v: :math:`(S, H, O)` where S is sparsed dimension, H is the number of
            heads, and O is output dimension.
        - index: :math:`(S)` where S is sparsed dimension.
        - dim_size: must be :math:`(B \times Nt)`

        - Output: attention values have shape :math:`(B, Nt, E)`; attention
            weights have shape :math:`(S, H)`
    """
    dim = 0
    B, H, E = q.shape
    N = maybe_num_nodes(index, dim_size)
    # scores
    alpha = (q * k).sum(dim=-1) / math.sqrt(E)
    alpha = sparse_softmax(alpha, index, num_nodes=N, dim=dim)
    if dropout_p > 0.0:
        alpha = F.dropout(alpha, p=dropout_p)
    v *= alpha.view(-1, H, 1)
    # out
    out = torch.zeros((N, H, v.size(2)), dtype=v.dtype, device=v.device)
    add_index = broadcast(index, v, dim)
    out.scatter_add_(dim, add_index, v)
    return out, alpha


from __future__ import print_function, division, absolute_import

import os
import copy
import json
from collections import OrderedDict
from warnings import warn
from six import iteritems, string_types

from sympy import Basic, sympify, Symbol
from numpy import bool_, float_
from jsonschema import validate, ValidationError
import cobra

import cobrame
from cobrame.util import me_model_interface, mu

try:
    # If cannot import SymbolicParameter, assume using cobrapy
    # versions <= 0.5.11
    from optlang.interface import SymbolicParameter
except ImportError:
    from cobra.io.json import metabolite_from_dict, save_json_model
else:
    from cobra.io.json import save_json_model
    from cobra.io.dict import metabolite_from_dict

mu_temp = Symbol('mu')

cur_dir = os.path.dirname(os.path.abspath(__file__))

def save_json_me(me0, file_name):
    """
    Save a stripped-down JSON version of the ME-model. This will exclude all of
    ME-Model information except the reaction stoichiometry information and the
    reaction bounds. Saving/loading a model in this format will thus occur much
    quicker, but limit the ability to edit the model and use most of its
    features.

    :param :class:`~cobrame.core.model.MEModel` me0:
        A full ME-model

    :param str or file-like object file_name:
        Filename of the JSON output

    :returns JSON-object:
        Stripped-down JSON representation of full ME-model
    """
    me = copy.deepcopy(me0)

    for rxn in me.reactions:
        for met in rxn.metabolites:
            s = rxn._metabolites[met]
            if isinstance(s, Basic):
                rxn._metabolites[met] = str(s)
        if isinstance(rxn.lower_bound, Basic):
            rxn.lower_bound = str(rxn.lower_bound)
        if isinstance(rxn.upper_bound, Basic):
            rxn.upper_bound = str(rxn.upper_bound)

    for met in me.metabolites:
        if isinstance(met._bound, Basic):
            met._bound = str(met._bound)

    save_json_model(me, file_name)

def get_sympy_expression(value):
    """
    Return sympy expression from json string using sympify

    mu is assumed to be positive but using sympify does not apply this
    assumption"""

    expression_value = sympify(value)
    return expression_value.subs(mu_temp, mu)

def get_numeric_from_string(string):
    try:
        return float(string)
    except ValueError:
        return get_sympy_expression(string)

def load_json_me(file_name):
    """
    Load a stripped-down JSON version of the ME-model. This will exclude all of
    ME-Model information except the reaction stoichiometry information and the
    reaction bounds. Saving/loading a model in this format will thus occur much
    quicker, but limit the ability to edit the model and use most of its
    features.

    :param str or file-like object file_name:
        Filename of the JSON ME-model

    Returns
    -------
    :class:`cobra.core.model.Model`
        COBRA Model representation of the ME-model. This will not include
        all of the functionality of a :class:`~cobrame.core.model.MEModel` but
        will solve identically compared to the full model.
    """
    if isinstance(file_name, string_types):
        with open(file_name, 'r') as f:
            obj = json.load(f)
    else:
        obj = file_name

    model = cobra.Model()

    # If cannot import SymbolicParameter, assume using cobrapy
    # versions <= 0.5.11. If versions >= 0.8.0 are used, a ME-model interface
    # must be assigned as the solver interface
    try:
        from optlang.interface import SymbolicParameter
    except ImportError:
        pass
    else:
        model.solver = me_model_interface

    default_reactions = [i.id for i in model.reactions]

    for k, v in iteritems(obj):
        if k in {'id', 'name'}:
            setattr(model, k, v)

    def _reaction_from_dict(reaction, model):
        new_reaction = cobra.Reaction()
        for k, v in iteritems(reaction):
            if k in {'objective_coefficient', 'reversibility', 'reaction'}:
                continue
            elif k == 'metabolites':
                new_reaction.add_metabolites(OrderedDict(
                    (model.metabolites.get_by_id(str(met)),
                     get_numeric_from_string(coeff))
                    for met, coeff in iteritems(v)))
            elif k in {'upper_bound', 'lower_bound'}:
                v = get_numeric_from_string(v)
                setattr(new_reaction, k, v)
            else:
                setattr(new_reaction, k, v)
        return new_reaction

    model.add_metabolites(
        [metabolite_from_dict(metabolite) for metabolite in obj['metabolites']]
    )

    new_reactions = [
        _reaction_from_dict(reaction, model) for reaction in obj['reactions']]

    model.remove_reactions(default_reactions)
    model.add_reactions(new_reactions)

    return model

# -----------------------------------------------------------------------------
# Functions below here facilitate json dumping/loading of full ME-models with
# all process_data/reaction info intact.
_REQUIRED_REACTION_ATTRIBUTES = {"id", "name", "metabolites", "lower_bound",
                                 "upper_bound", "objective_coefficient",
                                 "variable_kind"}

# Reaction types can have different attributes
_REACTION_TYPE_DEPENDENCIES = \
    {'MetabolicReaction': ['complex_data',
                           'stoichiometric_data',
                           'keff', 'reverse'],
     'ComplexFormation': ['_complex_id',
                          'complex_data_id'],
     'PostTranslationReaction':
         ['posttranslation_data'],
     'TranscriptionReaction': ['transcription_data'],
     'GenericFormationReaction': [],
     'MEReaction': [],
     'SummaryVariable': [],
     'TranslationReaction': ['translation_data'],
     'tRNAChargingReaction': ['tRNA_data']}

_REQUIRED_PROCESS_DATA_ATTRIBUTES = {"id"}

# Process data types have different attributes
_PROCESS_DATA_TYPE_DEPENDENCIES = \
    {'StoichiometricData': ['_stoichiometry', 'lower_bound', 'upper_bound',
                            'subreactions'],

     'ComplexData': ['stoichiometry', 'complex_id', 'subreactions'],

     'TranscriptionData': ['subreactions', 'nucleotide_sequence',
                           'RNA_products', 'RNA_polymerase'],

     'TranslationData': ['subreactions', 'nucleotide_sequence', 'mRNA',
                         'protein'],

     'tRNAData': ['subreactions', 'codon', 'RNA', 'amino_acid',
                  'synthetase', 'synthetase_keff'],

     'TranslocationData': ['enzyme_dict', 'stoichiometry', 'keff',
                           'length_dependent_energy'],

     'PostTranslationData': ['processed_protein_id', 'unprocessed_protein_id',
                             'propensity_scaling', 'aggregation_propensity',
                             'translocation', 'subreactions', 'surface_area',
                             'keq_folding', 'k_folding', 'biomass_type',
                             'translocation_multipliers'],

     'SubreactionData': ['stoichiometry', 'enzyme', 'keff',
                         'element_contribution'],

     'GenericData': ['component_list']
     }

_REQUIRED_METABOLITE_ATTRIBUTES = {"id", "name", "formula"}

_OPTIONAL_METABOLITE_ATTRIBUTES = {"charge", "formula", "compartment",
                                   "_bound", "_constraint_sense"}

# Some metabolite types require additional attributes
_METABOLITE_TYPE_DEPENDENCIES = \
    {'TranscribedGene': ['left_pos', 'right_pos', 'strand', 'RNA_type',
                         'nucleotide_sequence'],
     'ProcessedProtein': ['unprocessed_protein_id']
     }

def get_schema():
    with open(os.path.join(cur_dir, 'JSONSCHEMA'), 'r') as f:
        return json.load(f)

def _fix_type(value):
    """convert possible types to str, float, and bool"""
    # Because numpy floats can not be pickled to json
    if isinstance(value, string_types):
        return str(value)
    if isinstance(value, float_):
        return float(value)
    if isinstance(value, bool_):
        return bool(value)
    if isinstance(value, set):
        return list(value)
    if isinstance(value, Basic):
        return str(value)
    if hasattr(value, 'id'):
        return str(value.id)
    # if value is None:
    #     return ''
    return value

def _reaction_to_dict(reaction):
    new_reaction = {key: _fix_type(getattr(reaction, key))
                    for key in _REQUIRED_REACTION_ATTRIBUTES
                    if key != 'metabolites'}

    reaction_type = reaction.__class__.__name__
    new_reaction['reaction_type'] = {}
    new_reaction['reaction_type'][reaction_type] = {}

    for attribute in _REACTION_TYPE_DEPENDENCIES.get(reaction_type, []):
        reaction_attribute = getattr(reaction, attribute)

        new_reaction['reaction_type'][reaction_type][attribute] = \
            _fix_type(reaction_attribute)

    # Add metabolites
    new_reaction['metabolites'] = {}
    for met, value in reaction.metabolites.items():
        new_reaction['metabolites'][met.id] = _fix_type(value)

    return new_reaction

def _process_data_to_dict(data):
    process_data_type = data.__class__.__name__

    new_data = {key: _fix_type(getattr(data, key))
                for key in _REQUIRED_PROCESS_DATA_ATTRIBUTES}

    new_data['process_data_type'] = {}
    new_data['process_data_type'][process_data_type] = {}
    new_process_data_type_dict = \
        new_data['process_data_type'][process_data_type]

    special_list = ['subreactions', 'stoichiometry', 'enzyme_dict',
                    'surface_area', 'keq_folding' 'k_folding']

    for attribute in _PROCESS_DATA_TYPE_DEPENDENCIES[process_data_type]:
        if attribute not in special_list:
            data_attribute = getattr(data, attribute)

            new_process_data_type_dict[attribute] = _fix_type(data_attribute)

        elif attribute == 'enzyme_dict':
            new_process_data_type_dict[attribute] = {}
            for cplx, values in getattr(data, attribute).items():
                new_process_data_type_dict[attribute][cplx] = {}
                for property, value in values.items():
                    new_process_data_type_dict[attribute][cplx][property] = \
                        _fix_type(value)
        else:
            new_process_data_type_dict[attribute] = {}
            for metabolite, coefficient in getattr(data, attribute).items():
                new_process_data_type_dict[attribute][metabolite] = \
                    _fix_type(coefficient)

    return new_data

def _metabolite_to_dict(metabolite):

    metabolite_type = metabolite.__class__.__name__

    new_metabolite = {key: _fix_type(getattr(metabolite, key))
                      for key in _REQUIRED_METABOLITE_ATTRIBUTES}

    # Som metabolites require additional information to construct working
    # ME-model
    new_metabolite['metabolite_type'] = {}
    new_metabolite['metabolite_type'][metabolite_type] = {}
    for attribute in _METABOLITE_TYPE_DEPENDENCIES.get(metabolite_type, []):
        metabolite_attribute = getattr(metabolite, attribute)
        new_metabolite['metabolite_type'][metabolite_type][attribute] = \
            metabolite_attribute

    return new_metabolite

def get_attribute_array(dictlist, type):
    if type == 'reaction':
        return [_reaction_to_dict(reaction) for reaction in dictlist]
    elif type == 'process_data':
        return [_process_data_to_dict(data) for data in dictlist]
    elif type == 'metabolite':
        return [_metabolite_to_dict(metabolite) for metabolite in dictlist]
    else:
        raise TypeError('Type must be reaction, process_data or metabolite')

def get_global_info_dict(global_info):
    new_global_info = {}
    for key, value in global_info.items():
        if type(value) != dict:
            new_global_info[key] = _fix_type(value)
        else:
            new_global_info[key] = value
    return new_global_info

def _to_dict(model):

    obj = dict(
        reactions=get_attribute_array(model.reactions, 'reaction'),
        process_data=get_attribute_array(model.process_data,
                                         'process_data'),
        metabolites=get_attribute_array(model.metabolites, 'metabolite'),
        global_info=get_global_info_dict(model.global_info)
    )

    return obj

def save_full_me_model_json(model, file_name):
    """
    Save a full JSON version of the ME-model. Saving/loading a model in this
    format can then be loaded to return a ME-model identical to the one saved.

    :param :class:`~cobrame.core.MEModel.MEModel` model:
        A full ME-model

    :param str or file-like object file_name:
        Filename of the JSON output

    :returns JSON-object:
        Full JSON representation of full ME-model
    """

    should_close = False
    if isinstance(file_name, string_types):
        file_name = open(file_name, 'w')
        should_close = True

    json.dump(_to_dict(model), file_name)

    if should_close:
        file_name.close()

def add_metabolite_from_dict(model, metabolite_info):
    """
    Builds metabolite instances defined in dictionary, then add it to the
    ME-model being constructed.

    ProcessedProteins require additional information
    """

    metabolite_type_dict = metabolite_info['metabolite_type']
    if len(metabolite_type_dict) != 1:
        raise Exception('Only 1 metabolite_type in valid json')

    metabolite_type = list(metabolite_type_dict.keys())[0]

    # ProcessedProtein types require their unprocessed protein id as well
    if metabolite_type == 'ProcessedProtein':
        unprocessed_id = \
            metabolite_type_dict['ProcessedProtein']['unprocessed_protein_id']

        metabolite_obj = \
            getattr(cobrame, metabolite_type)(metabolite_info['id'],
                                              unprocessed_id)

    elif metabolite_type == 'TranscribedGene':
        rna_type = metabolite_type_dict['TranscribedGene']['RNA_type']
        nucleotide_sequence = \
            metabolite_type_dict['TranscribedGene']['nucleotide_sequence']
        metabolite_obj = \
            getattr(cobrame, metabolite_type)(metabolite_info['id'],
                                              rna_type, nucleotide_sequence)
    else:
        metabolite_obj = \
            getattr(cobrame, metabolite_type)(metabolite_info['id'])

    for attribute in _REQUIRED_METABOLITE_ATTRIBUTES:
        setattr(metabolite_obj, attribute, metabolite_info[attribute])

    for attribute in _METABOLITE_TYPE_DEPENDENCIES.get(metabolite_type, []):
        value = metabolite_type_dict[metabolite_type][attribute]
        setattr(metabolite_obj, attribute, value)

    model.add_metabolites([metabolite_obj])

def add_process_data_from_dict(model, process_data_dict):
    """
    Builds process_data instances defined in dictionary, then add it to the
    ME-model being constructed.

    Most classes of process_data only require an id and model to initiate them,
    but TranslationData, tRNAData, PostTranslationData and GenericData require
    additional inputs.

    """

    # Create process data instances. Handel certain types individually
    id = process_data_dict['id']
    process_data_type_dict = process_data_dict['process_data_type']
    if len(process_data_type_dict) == 1:
        process_data_type, process_data_info = process_data_type_dict.popitem()
    else:
        print(process_data_type_dict, len(process_data_type_dict))
        raise Exception('Only 1 reaction_type in valid json')

    if process_data_type == 'TranslationData':
        mrna = process_data_info['mRNA']
        protein = process_data_info['protein']
        process_data = \
            getattr(cobrame, process_data_type)(id, model, mrna, protein)
    elif process_data_type == 'tRNAData':
        amino_acid = process_data_info['amino_acid']
        rna = process_data_info['RNA']
        codon = process_data_info['codon']
        process_data = \
            getattr(cobrame, process_data_type)(id, model, amino_acid, rna,
                                                codon)
    elif process_data_type == 'PostTranslationData':
        processed_protein_id = process_data_info['processed_protein_id']
        unprocessed_protein_id = process_data_info['unprocessed_protein_id']
        process_data = \
            getattr(cobrame, process_data_type)(id, model,
                                                processed_protein_id,
                                                unprocessed_protein_id)
    elif process_data_type == 'GenericData':
        component_list = process_data_info['component_list']
        process_data = \
            getattr(cobrame, process_data_type)(id, model, component_list)
        # Create reaction from generic process data
        process_data.create_reactions()
    else:
        process_data = getattr(cobrame, process_data_type)(id, model)

    # Set all of the required attributes using information in info dictionary
    for attribute in _REQUIRED_PROCESS_DATA_ATTRIBUTES:
        setattr(process_data, attribute, process_data_dict[attribute])

    # Some attributes depend on process data type. Set those here.
    for attribute in _PROCESS_DATA_TYPE_DEPENDENCIES.get(process_data_type,
                                                         []):
        value = process_data_info[attribute]
        try:
            setattr(process_data, attribute, value)
        except AttributeError:
            # set to the hidden attribute instead
            setattr(process_data, '_' + attribute, value)

def add_reaction_from_dict(model, reaction_info):
    """
    Builds reaction instances defined in dictionary, then add it to the
    ME-model being constructed.

    """
    reaction_type_dict = reaction_info['reaction_type']

    if len(reaction_type_dict) == 1:
        reaction_type = list(reaction_type_dict.keys())[0]
        reaction_obj = getattr(cobrame, reaction_type)(reaction_info['id'])
    else:
        raise Exception('Only 1 reaction_type in valid json')

    for attribute in _REQUIRED_REACTION_ATTRIBUTES:
        # Metabolites are added to reactions using their update function,
        # skip setting metabolite stoichiometries here
        if attribute == 'metabolites':
            continue

        # upper and lower bounds may contain mu values. Handle that here
        value = reaction_info[attribute]
        if attribute in ['upper_bound', 'lower_bound']:
            value = get_sympy_expression(value)
        setattr(reaction_obj, attribute, value)

    # Some reactions are added to model when ME-models are initialized
    try:
        model.add_reactions([reaction_obj])
    except Exception:
        reaction_obj = model.reactions.get_by_id(reaction_obj.id)
        if reaction_type not in ['SummaryVariable',
                                 'GenericFormationReaction']:
            warn('Reaction (%s) already in model' % reaction_obj.id)

    # These reactions types do not have update functions and need their
    # stoichiometries set explicitly .
    if reaction_type in ['SummaryVariable', 'MEReaction']:
        for key, value in reaction_info['metabolites'].items():
            reaction_obj.add_metabolites({key: get_sympy_expression(value)},
                                         combine=False)

    for attribute in _REACTION_TYPE_DEPENDENCIES.get(reaction_type, []):
        # Spontaneous reactions do no require complex_data
        if attribute == 'complex_data' and 'SPONT' in reaction_obj.id:
            continue

        value = reaction_type_dict[reaction_type][attribute]
        setattr(reaction_obj, attribute, value)

    if hasattr(reaction_obj, 'update'):
        reaction_obj.update()

def full_me_model_from_dict(obj):
    """
    Validate and load JSON representation of the ME-model. This will return
    a full :class:`~cobrame.core.model.MEModel` object identical to the
    one saved.

    :param str or file-like object obj:
        JSON-serialized ME-model

    :returns :class:`~cobrame.core.model.MEModel`:
        Full COBRAme ME-model
    """

    try:
        validate(obj, get_schema())
    except ValidationError:
        raise Exception('Must pass valid ME-model json file')

    model = cobrame.MEModel()

    for k, v in iteritems(obj):
        if k in {'id', 'name', 'global_info'}:
            setattr(model, k, v)

    for metabolite in obj['metabolites']:
        add_metabolite_from_dict(model, metabolite)

    for process_data in obj['process_data']:
        add_process_data_from_dict(model, process_data)

    for reaction in obj['reactions']:
        add_reaction_from_dict(model, reaction)

    model.update()

    return model

def load_full_me_model_json(file_name):

    with open(file_name, 'r') as f:
        model_dict = json.load(f)

    return full_me_model_from_dict(model_dict)