"""compare data from different sources to see what has changed,
typically previous and updated results"""

import lib_leiden
from matplotlib import pyplot as plt
from scipy import interpolate,integrate
import numpy as np

## directories to compare
directory0 = 'td'
directory1 = '../'
fractional_change_threshold_for_notification = 3e-2
one_species_to_plot = 'CO'

## compare data in photo_rates.csv
d0 = lib_leiden.txt_to_dict(f'{directory0}/photo_rates.csv',labels_commented=False,delimiter=',')
d1 = lib_leiden.txt_to_dict(f'{directory1}/photo_rates.csv',labels_commented=False,delimiter=',')
assert set(d0.keys())==set(d1.keys()),'Keys differ'
for i0 in range(len(d0['species'])):
    ## check matching species/products
    i = (d1['species']==d0['species'][i0])&(d1['products']==d0['products'][i0])
    if sum(i)==0:
        print(f"photo_rates.csv: {d0['species'][i0]:10} {d0['products'][i0]:20} not in {directory1}")
        continue
    if sum(i)>1:
        print("photo_rates.csv: double matches?!?")
        continue
    i1 = int(np.where(i)[0])
    ## check differing rates
    for key in d0.keys():
        if key in ('species','products'):
            continue
        if (change:=(d0[key][i0]-d1[key][i1])/d0[key][i0]) > fractional_change_threshold_for_notification:
            print(f'photo_rates.csv: {d0["species"][i0]:10} {d0["products"][i0]:20} {key:10} large fractional change: {change:0.4f}')
for i1 in range(len(d1['species'])):
    ## check matching species/products
    i = (d0['species']==d1['species'][i1])&(d0['products']==d1['products'][i1])
    if sum(i)==0:
        print(f"photo_rates.csv: {d1['species'][i1]:10} {d1['products'][i1]:20} not in {directory0}")
        continue

## compare in comic_ray_rates.csv
d0 = lib_leiden.txt_to_dict(f'{directory0}/cosmic_ray_rates.csv',labels_commented=False,delimiter=',')
d1 = lib_leiden.txt_to_dict(f'{directory1}/cosmic_ray_rates.csv',labels_commented=False,delimiter=',')
assert set(d0.keys())==set(d1.keys()),'Keys differ'
for i0 in range(len(d0['species'])):
    ## check matching species/products
    i = (d1['species']==d0['species'][i0])&(d1['products']==d0['products'][i0])
    if sum(i)==0:
        print(f"cosmic_ray_rates.csv: {d0['species'][i0]:10} {d0['products'][i0]:20} not in {directory1}")
        continue
    if sum(i)>1:
        print("cosmic_ray_rates.csv: double matches?!?")
        continue
    i1 = int(np.where(i)[0])
    if (change:=(d0['rate'][i0]-d1['rate'][i1])/d0['rate'][i0]) > fractional_change_threshold_for_notification:
        print(f'cosmic_ray_rates.csv: {d0["species"][i0]:10} {d0["products"][i0]:20} large fractional change: {change:0.4f}')
for i1 in range(len(d1['species'])):
    ## check matching species/products
    i = (d0['species']==d1['species'][i1])&(d0['products']==d1['products'][i1])
    if sum(i)==0:
        print(f"cosmic_ray_rates.csv: {d1['species'][i1]:10} {d1['products'][i1]:20} not in {directory0}")
        continue

## compare data in gamma_factors.csv
d0 = lib_leiden.txt_to_dict(f'{directory0}/gamma_factors.csv',labels_commented=False,delimiter=',')
d1 = lib_leiden.txt_to_dict(f'{directory1}/gamma_factors.csv',labels_commented=False,delimiter=',')
assert set(d0.keys())==set(d1.keys()),'Keys differ'
## find common data
hash0 = [hash((d0['species'][i],d0['products'][i],d0['radiation_field'][i],d0['grains_type'][i])) for i in range(len(d0))]
hash1 = [hash((d1['species'][i],d1['products'][i],d1['radiation_field'][i],d1['grains_type'][i])) for i in range(len(d1))]
i0,i1 = lib_leiden.common(hash0,hash1)
## check for missing data
if len(i0)<len(d0):
    j = lib_leiden.isin(d0[i0],d0)
    for k in my.find(~j):
        print(f'gamma_factors.csv: {d0["species"][k]:10} {d0["products"][k]:20} {d0["radiation_field"][k]:20} {d0["grains_type"][k]:20} some data not in d0 ')
if len(i1)<len(d1):
    j = lib_leiden.isin(d1[i1],d1)
    for k in my.find(~j):
        print(f'gamma_factors.csv: {d1["species"][k]:11} {d1["products"][k]:20} {d1["radiation_field"][k]:20} {d1["grains_type"][k]:20} some data not in d1 ')
## check data valeus
for j0,j1 in zip(i0,i1):
    for key in ('gamma_E2','gamma_exp','gamma_exp_fit','k0_exp_fit/k0'):
        if (change:=(d0[key][j0]-d1[key][j1])/d0[key][j0]) > fractional_change_threshold_for_notification:
            print(f'gamma_factors.csv: {d0["species"][j0]:10} {d0["products"][j0]:20} {d0["radiation_field"][j0]:20} {d0["grains_type"][j0]:20} large fractional change: {change:10.4f}')

## add dust shielding function comparison here

## add line shielding function comparison here

## plot data for one particular species
if one_species_to_plot is not None:
    
    from anh import * # nonlocal dependency!

    species = one_species_to_plot
    
    ## compare h5 cross section
    fig = plt.figure(1)
    fig.clf()
    ax = fig.gca()
    x0 = lib_leiden.hdf5_to_dict(f'{directory0}/cross_sections/{species}/{species}.h5')
    x1 = lib_leiden.hdf5_to_dict(f'{directory1}/cross_sections/{species}/{species}.h5')
    ax.set_title('cross section full resolution')
    print('integrated cross section full resolution')
    print(format('products','18'),format(directory0,'>15'),format(directory1,'>15'),)
    for ikey,key in enumerate(x0.keys()):
        if key in ('wavelength','README'): continue
        ax.plot(x0['wavelength'],x0[key],label=f'{species:15} {key:15} {directory0}',color=lib_leiden.newcolor(ikey),linestyle='-')
        ax.plot(x1['wavelength'],x1[key],label=f'{species:15} {key:15} {directory1}',color=lib_leiden.newcolor(ikey),linestyle=':')
        print(format(key,'18'),
              format(integrate.trapz(x0[key],x0['wavelength']),'15e'),
              format(integrate.trapz(x1[key],x1['wavelength']),'15e'),)
    lib_leiden.legend(show_style=True)
    ax.set_yscale('log')

    ## compare 0.1nm cross section
    fig = plt.figure(2)
    fig.clf()
    ax = fig.gca()
    x0 = lib_leiden.txt_to_dict(f'{directory0}/cross_sections/{species}/{species}_0.1nm.txt')
    x1 = lib_leiden.txt_to_dict(f'{directory1}/cross_sections/{species}/{species}_0.1nm.txt')
    ax.set_title('cross section 0.1nm')
    print('integrated cross section full resolution')
    print(format('products','18'),format(directory0,'>15'),format(directory1,'>15'),)
    for ikey,key in enumerate(x0.keys()):
        if key in ('wavelength','README'): continue
        ax.plot(x0['wavelength'],x0[key],label=f'{species:15} {key:15} {directory0}',color=lib_leiden.newcolor(ikey),linestyle='-')
        ax.plot(x1['wavelength'],x1[key],label=f'{species:15} {key:15} {directory1}',color=lib_leiden.newcolor(ikey),linestyle=':')
        print(format(key,'18'),
              format(integrate.trapz(x0[key],x0['wavelength']),'15e'),
              format(integrate.trapz(x1[key],x1['wavelength']),'15e'),)
    lib_leiden.legend(show_style=True)
    ax.set_yscale('log')

    ## compare photo rates
    d0 = file_to_dynamic_recarray(f'{directory0}/photo_rates.csv')
    d0.limit_to_matches(species=species)
    d1 = file_to_dynamic_recarray(f'{directory1}/photo_rates.csv')
    d0,d1 = get_common(d0,d1,'species','products')
    # print( d0)
    my.fig(3)
    keys = list(d0.keys())
    keys.remove('species')
    keys.remove('products')
    # print( d0[keys])
    for iproducts,products in enumerate(d0.unique('products')):
        subplot(iproducts)
        title(f'photo rates {species} {products}')
        t0 = d0.matches(products=products)
        t1 = d1.matches(products=products)
        x = np.arange(len(keys))
        plt.bar(
            x,
            [float(t0[key]) for key in keys],
            align='edge',
            width=-0.4,
            label=f'{directory0}',color=my.newcolor(0))
        plt.bar(
            x,
            [float(t1[key]) for key in keys],
            align='edge',
            width=0.4,
            label=f'{directory1}',color=my.newcolor(1),)
        # plt.xticks(x,keys)
        my.set_tick_labels_text(keys)
        plt.legend()
        yscale('log')

    ## compare cosmic rates
    d0 = file_to_dynamic_recarray(f'{directory0}/cosmic_ray_rates.csv')
    d0.limit_to_matches(species=species)
    d1 = file_to_dynamic_recarray(f'{directory1}/cosmic_ray_rates.csv')
    d0,d1 = get_common(d0,d1,'species','products')
    my.fig(4)
    title(f'cosmic ray rates')
    x = np.arange(len(d0))
    plt.bar(x, d0['rate'], align='edge', width=-0.4, label=f'{directory0}',color=my.newcolor(0))
    plt.bar(x, d1['rate'], align='edge', width=0.4, label=f'{directory1}',color=my.newcolor(1))
    yscale('log')
    my.set_tick_labels_text([t['species']+'/'+t['products'] for t in d0])

    ## gamma factors
    d0 = file_to_dynamic_recarray(f'{directory0}/gamma_factors.csv')
    d0.limit_to_matches(species=species)
    d1 = file_to_dynamic_recarray(f'{directory1}/gamma_factors.csv')
    d0,d1 = get_common(d0,d1,'species','products','radiation_field','grains_type')
    for isubplot,(species,products,grains_type) in enumerate(d0.unique_combinations('species','products','grains_type')):
        my.fig(5+isubplot)
        title(f'gamma factors {species} {products} {grains_type}')
        keys = ('gamma_E2' , 'gamma_exp' , 'gamma_exp_fit' , 'k0_exp_fit/k0')
        radiation_fields = d0.unique('radiation_field')
        x = range(len(keys))
        xlabels = []
        for ix,(key,radiation_field) in enumerate(itertools.product(keys,radiation_fields)):
            i = d0.match(species=species,products=products,grains_type=grains_type,radiation_field=radiation_field)
            xlabels.append(f'{key} {radiation_field}')
            if sum(i)==0: continue
            plt.bar(ix, [float(d0[key][i]) for key in keys], align='edge', width=-0.4, label=f'{directory0}',color=my.newcolor(0))
            plt.bar(ix, [float(d1[key][i]) for key in keys], align='edge', width=+0.4, label=f'{directory1}',color=my.newcolor(1))
        yscale('log')
        my.set_tick_labels_text(xlabels)

    ## could add shielding function comparison here

    plt.show()