#!/usr/bin/env python3

"""Look for problems in the cross sections."""

import os,re
import lib_leiden
import numpy as np
from scipy import integrate

# cross_sections_directory = '../cross_sections'
cross_sections_directory = 'td/cross_sections'

for dirpath,dirnames,filenames in os.walk(cross_sections_directory):

    ## unusual subdirectories to skip
    if dirpath in (
            cross_sections_directory,
            f'{cross_sections_directory}/N2/isotopologues',
    ):
        continue

    species = re.sub(r'.*\/(.*)',r'\1',dirpath)

    print(f'testing: {species}')

    ## check all files present
    for t in (
            f'{species}.txt',
            f'{species}.pdf',
            f'{species}.h5',
            f'{species}.png',
            f'{species}_0.1nm.txt',
    ):
        if t not in filenames:
            print(f'{species}: missing {t}')

    ## load h5 cross section and test things
    try:

        filename = f'{dirpath}/{species}.h5'
        if not os.path.exists(filename):
            print(f'{species}: Cannot find cross section')
            print( dirpath)
            continue
        data = lib_leiden.hdf5_to_dict(filename)

        ## photoabsorption should always exist
        if 'photoabsorption' not in data:
            print(f'{species}: Photoabsorption cross section missing')
            continue
        
        ## look for negative cross sections
        if 'photoabsorption' in data and np.any(data['photoabsorption']<0):
            print(f'{species}: Photoabsorption cross section less than zero')
        if 'photodissociation' in data and np.any(data['photodissociation']<0):
            print(f'{species}: Photodissociation cross section less than zero')
        if 'photoionisation' in data and np.any(data['photoionisation']<0):
            print(f'{species}: Photoionisation cross section less than zero')
        
        ## look for photoionisation+photoionisation~=photabsorption
        t = x = data['photoabsorption']
        if 'photodissociation' in data:
            t -= data['photodissociation']
        if 'photoionisation' in data:
            t -= data['photoionisation']
        if np.any(np.abs(x-t)>1e-24):
            print(f'{species}: Photoionisation and photodissociation cross sections do no add up to photoabsorption')
        ## monotonic wavelength
        if np.any(np.diff(data['wavelength'])<=0):
            print(f'{species}: Wavelength scale not monotonically increasing')

    except Exception as err:

        print(f'error in species: {species}')
        raise err
        
    ## load different realisations of each cross section
    data_h5 = lib_leiden.hdf5_to_dict(f'{dirpath}/{species}.h5')
    data_txt = lib_leiden.txt_to_dict(f'{dirpath}/{species}.txt')
    data_01nm_txt = lib_leiden.txt_to_dict(f'{dirpath}/{species}_0.1nm.txt')

##     ## test beginning and end wavelengths and values are the same - NOT CURRENTLY TESTING PD AND PI FILES
##     for key in ('wavelength','photoabsorption','photodissociation','photoionisation'):
##         if key not in data_h5:
##             continue
##         if abs(data_h5[key][0]-data_txt[key][0])>1e-25:
##                 print(f'{species}: Beginning {key} of h5 and txt files do not match. {data_h5[key][0]} and {data_txt[key][0]}')
##         if abs(data_h5[key][-1]-data_txt[key][-1])>1e-25:
##                 print(f'{species}: Ending {key} of h5 and txt files do not match. {data_h5[key][-1]} and {data_txt[key][-1]}')
##         if abs(data_h5[key][0]-data_01nm_txt[key][0])>1e-25:
##                 print(f'{species}: Beginning {key} of h5 and 0.1nm_txt files do not match. {data_h5[key][0]} and {data_01nm_txt[key][0]}')
##         if abs(data_h5[key][-1]-data_01nm_txt[key][-1])>1e-25:
##                 print(f'{species}: Ending {key} of h5 and 0.1nm_txt files do not match. {data_h5[key][-1]} and {data_01nm_txt[key][-1]}')
        
    ## ensure integrated values are the same
    int_h5 = integrate.trapz(data_h5['photoabsorption'],data_h5['wavelength']) 
    int_txt = integrate.trapz(data_txt['photoabsorption'],data_txt['wavelength']) 
    int_01nm_txt = integrate.trapz(data_01nm_txt['photoabsorption'],data_01nm_txt['wavelength'])
    difference=abs(int_h5-int_txt)/int_h5
    if difference > 1e-5:
        print(f'{species}: Integrated cross sections h5 and txt files do not match. Difference = {difference}.')
        print(f'{species}: Integrated cross sections h5 and txt files do not match.')
    difference=abs(int_h5-int_01nm_txt)/int_h5
    if difference > 5e-3:
        print(f'{species}: Integrated cross sections h5 and 0.1nm_txt files do not match. Difference = {difference}.')