
#%% imports
import shutil
from turtle import pd
from pyhcal.calibrators import calibrator
from pyhcal.metrics import calculate_all_statistics,align_series
from pathlib import Path
import pyemu
import pandas as pd
#%% Hydrology Defaults TODO: Move to separate module
# Combined PEST defaults AND metadata for HSPF parameters
HSPF_PARAMETER_DEFAULTS = {
    # ===== PERLND HYDROLOGY (PWATER) =====
    'LZSN': {
        'partrans': 'log',
        'parchglim': 'factor',
        'parval1': 5.0,
        'parlbnd': 2.0,
        'parubnd': 15.0,
        'pargp': 'soil_storage',
        'scale': 1.0,
        'offset': 0.0,
        'dercom': 1,
        'units': 'inches',
        'description': 'Lower zone nominal storage',
        'module': 'PERLND',
        'typical_table': 'PWAT-PARM2',
        'activity': 'PWATER',
    },
    'INFILT': {
        'partrans': 'log',
        'parchglim': 'factor',
        'parval1': 0.2,
        'parlbnd': 0.01,
        'parubnd': 0.5,
        'pargp': 'infiltration',
        'scale': 1.0,
        'offset': 0.0,
        'dercom': 1,
        'units': 'in/hr',
        'description': 'Infiltration index',
        'module': 'PERLND',
        'typical_table': 'PWAT-PARM2',
        'activity': 'PWATER',
    },
    'UZSN': {
        'partrans': 'log',
        'parchglim': 'factor',
        'parval1': 0.5,
        'parlbnd': 0.1,
        'parubnd': 2.0,
        'pargp': 'soil_storage',
        'scale': 1.0,
        'offset': 0.0,
        'dercom': 1,
        'units': 'inches',
        'description': 'Upper zone nominal storage',
        'module': 'PERLND',
        'typical_table': 'PWAT-PARM2',
        'activity': 'PWATER',
    },
    'AGWRC': {
        'partrans': 'none',
        'parchglim': 'relative',
        'parval1': 0.98,
        'parlbnd': 0.85,
        'parubnd': 0.999,
        'pargp': 'groundwater',
        'scale': 1.0,
        'offset': 0.0,
        'dercom': 1,
        'units': '1/day',
        'description': 'Active groundwater recession rate',
        'module': 'PERLND',
        'typical_table': 'PWAT-PARM3',
        'activity': 'PWATER',
    },
    'BASETP': {
        'partrans': 'log',
        'parchglim': 'factor',
        'parval1': 0.05,
        'parlbnd': 0.0,
        'parubnd': 0.2,
        'pargp': 'evapotranspiration',
        'scale': 1.0,
        'offset': 0.0,
        'dercom': 1,
        'units': 'none',
        'description': 'Baseflow evapotranspiration parameter',
        'module': 'PERLND',
        'typical_table': 'PWAT-PARM3',
        'activity': 'PWATER',
    },
    'DEEPFR': {
        'partrans': 'none',
        'parchglim': 'relative',
        'parval1': 0.1,
        'parlbnd': 0.0,
        'parubnd': 0.5,
        'pargp': 'groundwater',
        'scale': 1.0,
        'offset': 0.0,
        'dercom': 1,
        'units': 'fraction',
        'description': 'Fraction of groundwater inflow to deep/inactive groundwater',
        'module': 'PERLND',
        'typical_table': 'PWAT-PARM3',
        'activity': 'PWATER',
    },
    'INTFW': {
        'partrans': 'log',
        'parchglim': 'factor',
        'parval1': 2.0,
        'parlbnd': 1.0,
        'parubnd': 10.0,
        'pargp': 'interflow',
        'scale': 1.0,
        'offset': 0.0,
        'dercom': 1,
        'units': 'none',
        'description': 'Interflow index',
        'module': 'PERLND',
        'typical_table': 'PWAT-PARM2',
        'activity': 'PWATER',
    },
    'IRC': {
        'partrans': 'none',
        'parchglim': 'relative',
        'parval1': 0.5,
        'parlbnd': 0.3,
        'parubnd': 0.85,
        'pargp': 'interflow',
        'scale': 1.0,
        'offset': 0.0,
        'dercom': 1,
        'units': '1/day',
        'description': 'Interflow recession parameter',
        'module': 'PERLND',
        'typical_table': 'PWAT-PARM2',
        'activity': 'PWATER',
    },
    'KVARY': {
        'partrans': 'log',
        'parchglim': 'factor',
        'parval1': 1.0,
        'parlbnd': 0.0,
        'parubnd': 5.0,
        'pargp': 'groundwater',
        'scale': 1.0,
        'offset': 0.0,
        'dercom': 1,
        'units': 'in-1',
        'description': 'Variable groundwater recession parameter',
        'module': 'PERLND',
        'typical_table': 'PWAT-PARM3',
        'activity': 'PWATER',
    },
    'AGWETP': {
        'partrans': 'none',
        'parchglim': 'relative',
        'parval1': 0.1,
        'parlbnd': 0.0,
        'parubnd': 0.3,
        'pargp': 'evapotranspiration',
        'scale': 1.0,
        'offset': 0.0,
        'dercom': 1,
        'units': 'none',
        'description': 'Active groundwater evapotranspiration parameter',
        'module': 'PERLND',
        'typical_table': 'PWAT-PARM3',
        'activity': 'PWATER',
    },
    'CEPSC': {
        'partrans': 'log',
        'parchglim': 'factor',
        'parval1': 0.2,
        'parlbnd': 0.05,
        'parubnd': 0.5,
        'pargp': 'interception',
        'scale': 1.0,
        'offset': 0.0,
        'dercom': 1,
        'units': 'inches',
        'description': 'Interception storage capacity',
        'module': 'PERLND',
        'typical_table': 'PWAT-PARM2',
        'activity': 'PWATER',
    },
    'LZETP': {
        'partrans': 'none',
        'parchglim': 'relative',
        'parval1': 0.7,
        'parlbnd': 0.1,
        'parubnd': 0.9,
        'pargp': 'evapotranspiration',
        'scale': 1.0,
        'offset': 0.0,
        'dercom': 1,
        'units': 'none',
        'description': 'Lower zone ET parameter',
        'module': 'PERLND',
        'typical_table': 'PWAT-PARM3',
        'activity': 'PWATER',
    }
}


#%% Select calibrator and model file to build pest project
folder_path = Path('C:/Users/mfratki/Documents/Projects/Tests')

model_name = 'Nemadji'
cal = calibrator(folder_path / model_name)
cal.pest_path = cal.project_path / 'pest'
cal.load_model(1)

# reach_ids = []
# for outlet in outlets.values():
#     reach_ids.extend(outlet['reach_ids'])
# reach_ids = list(set(reach_ids))

# cal.initialize(reach_ids)
# cal.uci.set_simulation_period(2015,2018)
# cal.run_model('pest')
#%% 
#Copy this ucefile and associated files to the pest folder
shutil.copy(cal.uci.filepath, cal.pest_path / cal.uci.filepath.name)
for file in cal.uci.table('FILES')['FILENAME'].values:
    # copy these files to pest location
    src = cal.model_path / file
    dst = cal.pest_path / file
    shutil.copy(src,dst)

org_d = cal.pest_path
tmp_d = cal.pest_path / 'tmp'
template_ws = tmp_d / 'template'
if tmp_d.exists():
    shutil.rmtree(tmp_d)
shutil.copytree(org_d,tmp_d)


uci_file = cal.uci.filepath.name
uci_tpl = cal.uci.filepath.name.replace('.uci', '.tpl')


#%% Parameter configuration
# parameters to be calibrated


parameters = [
            {'block': 'PERLND',
            'table': 'PWAT-PARM2',
            'table_id': 0,
            'hspf_parameter' : 'LZSN',
            'single_template' : False,
            'opnids' : None,
            'group' : '',
            'parname' : 'lzsn'
            },
            {'block': 'PERLND',
            'table': 'PWAT-PARM2',
            'table_id': 0,
            'hspf_parameter' : 'INFILT',
            'single_template' : False,
            'opnids' : None,
            'group' : '',
            'parname' : 'inflt'
            }
]
    
for parm in parameters:
    parm['pargp'] = f"{parm['hspf_parameter']}_{parm['group']}"

#%% Build HSPF parameter template and store parameter data for later
parameter_data = []
for parameter in parameters:
    parnme = cal.uci.add_parameter_template(parameter['block'], 
                                   parameter['table'], 
                                   parameter['table_id'], 
                                   parameter['hspf_parameter'],
                                   single_template=parameter['single_template'],
                                   opnids=parameter['opnids'],
                                   group_id=parameter['group'],
                                   parname=parameter['parname'])
    row = HSPF_PARAMETER_DEFAULTS.get(parameter['hspf_parameter']) | parameter
    rows = [row]*len(parnme)
    df = pd.DataFrame(rows)
    df['parnme'] = parnme
    df = df.set_index('parnme')
    parameter_data.append(df)

parameter_data = pd.concat(parameter_data)


#%% Observation Configuration calibration and sensitivity analysis

#Observations configuration
# Stations

observations = [{'name': 'stn1',     #  unique name for this observation location; used to link observed data to model output and parameter sensitivities
                'station_id': 'E05011002', #  ID for the station; used to link observed data to model output and parameter sensitivities; this should match the station ID used in the calibrator's get_observed_data method
                'station_origin': 'wiski', #  source of the observed data; used to link observed data to model output and parameter sensitivities; this should match the station origin used in the calibrator's get_observed_data method
                'station_name': 'Nemadji_Outlet', #optional descriptive name for the station; not used for linking data but can be helpful for organization and interpretation
                'station_weight': 1.0, #optional weight defaults to 1.0; used to weight the importance of this observation in the calibration objective function; higher values give more weight to this observation
                'reach_ids': [103], # required list of model reach IDs that correspond to this observation; used to link observed data to model output and parameter sensitivities; these should match the reach IDs used in the calibrator's get_simulated_data method
                'true_reach_id': 103, #optional true reach ID; if provided, this can be used to link observed data to model output and parameter sensitivities in cases where the station is not located at the outlet of a single reach or where the model grid does not align well with observation locations; this should match the reach ID used in the calibrator's get_simulated_data method for this station
                'model_name': model_name, # required model name; used to link observed data to model output and parameter sensitivities; this should match the model name used in the calibrator's get_simulated_data method
                'constituent': 'Q', # required constituent; used to link observed data to model output and parameter sensitivities; this should match the constituent name used in the calibrator's get_simulated_data method
                'time_step': 'D', #required time step; used to link observed data to model output and parameter sensitivities; this should match the time step used in the calibrator's get_simulated_data method
                'metrics': ['nse','lognse','bias','rmse','monthly_avg','exceedance']} #optional list of metrics to calculate for this observation; defaults to all available metrics;                ]
                ]


#                 },
#                 {'name': 'stn2',
#                 'station_id': 'S001-235',
#                 'station_origin': 'equis',
#                 'station_name': 'Nemadji_Outlet',
#                 'station_weight': 1.0,
#                 'reach_ids': [103],
#                 'true_reach_id': 103,
#                 'model_name': model_name,
#                 'constituent': 'Q',
#                 'time_step': 'D',
#                 'metrics': {'nse': 2.0,
#                             'lognse': 1.0,
#                             'bias': 1.0,
#                             'rmse': 1.0,
#                             'monthly_avg': 1.0,
#                             'exceedance': 1.0}
#                         }

#                 {'name': 'stn3',
#                 'station_id': 'S005-115',
#                 'station_origin': 'equis',
#                 'station_name': 'Nemadji_Outlet',
#                 'station_weight': 1.0,
#                 'reach_ids': [103],
#                 'true_reach_id': 103,
#                 'model_name': model_name,
#                 'constituent': 'Q',
#                 'time_step': 'D',
#                 'metrics': {'nse': 2.0,
#                             'lognse': 1.0,
#                             'bias': 1.0,
#                             'rmse': 1.0,
#                             'monthly_avg': 1.0,
#                             'exceedance': 1.0}
#                         }
# ]
# Internal constraints

#%% Internal constraints configuration
constraints = None


#%% Store initial model run to generate observations for sensitivty analysis

cal.model.outputs.set_output_folder(tmp_d)
reach_output_files = []
for observation in observations:
    reach_output_files.append(cal.model.outputs.write_outlet_output(
                                                     name = observation['name'], 
                                                     reach_ids = observation['reach_ids'], 
                                                     time_step = observation['time_step']))
catchment_output_file = cal.model.outputs.write_catchment_loading_output()

output_files = reach_output_files + [catchment_output_file]
#%%
# observation_df = []
# metrics_df = []
# observed_df = []
# for observation in observations:
#     constituent = observation['constituent']
#     reach_ids = observation['reach_ids']
#     time_step = observation['time_step']
#     selected_metrics = observation['metrics']
    
#     simulated = cal.model.hbns.get_reach_constituent(constituent,reach_ids,time_step)[0]
#     observed = cal.get_observed_data(observation['station_id'], constituent, time_step)
    
#     obs_aligned, sim_aligned = align_series(observed,simulated)
#     metrics_data = calculate_all_statistics(obs_aligned,sim_aligned)
#     metrics_data = pd.DataFrame([metrics_data]).T
#     metrics_data = metrics_data.dropna()
    
#     metrics_data.reset_index(inplace=True)
#     metrics_data.columns = ['metric','value']
#     metrics_data['name'] = observation['name']
#     metrics_data['constituent'] = observation['constituent']

#     sim_aligned = sim_aligned.reset_index()
#     sim_aligned.columns = ['datetime','value']
#     sim_aligned['name'] = observation['name']
#     sim_aligned['constituent'] = observation['constituent']
    
#     obs_aligned = obs_aligned.reset_index()
#     obs_aligned.columns = ['datetime','value']
#     obs_aligned['name'] = observation['name']
#     obs_aligned['constituent'] = observation['constituent']

#     observation_df.append(sim_aligned)
#     metrics_df.append(metrics_data)
#     observed_df.append(obs_aligned)

# observation_df = pd.concat(observation_df)
# metrics_df = pd.concat(metrics_df)
# observed_df = pd.concat(observed_df)

# metrics_output_path = tmp_d / 'metrics.csv'
# observations_output_path = tmp_d / 'observations.csv'
# observed_output_path = tmp_d / 'observed.csv'



# observed_df.to_csv(observed_output_path,index=False)
# metrics_df.to_csv(metrics_output_path,index=False)
# observation_df.to_csv(observations_output_path,index=False)

#%%
#############################
### Build PEST project ######
#############################

# #%% # instantiate PstFrom
pf = pyemu.utils.PstFrom(original_d=tmp_d, # where the model is stored
                            new_d=template_ws, # the PEST template folder
                            remove_existing=True, # ensures a clean start
                            longnames=True, # set False if using PEST/PEST_HP
                            zero_based=False, # does the MODEL use zero based indices? For example, MODFLOW does NOT
                            echo=True) # to stop PstFrom from writing lots of information to the notebook; experiment by setting it as True to see the difference; useful for troubleshooting


#%% Build uci template and store parameter data for later
cal.uci.write_tpl(new_tpl_path = template_ws / uci_tpl)


#%% Create PEST observation instruction files
for reach_output in reach_output_files:
    pf.add_observations(reach_output.name, 
                        insfile= f"{reach_output.name}.ins", 
                        index_cols=['datetime','name','constituent'],
                        use_cols=['value'],
                        prefix=reach_output.stem)

pf.add_observations(catchment_output_file.name,
                    insfile= f"{catchment_output_file.name}.ins", 
                    index_cols=['TVOLNO','constituent'],
                    use_cols=['loading_rate'],
                    prefix=catchment_output_file.stem)
    
    
# hds_df = pf.add_observations('observations.csv', # the model output file to read
#                                 insfile='observations.csv.ins', #optional, the instruction file name
#                                 index_cols= ['datetime','name','constituent'], #column header to use as index; can also use column number (zero-based) instead of the header name
#                                 use_cols=['observed'], #names of columns that include observation values; can also use column number (zero-based) instead of the header name
#                                 prefix="timeseries") #prefix to all observation names; choose something logical and easy to find. We use it later on to select observations

# stats_df = pf.add_observations('statistics.csv', # the model output file to read
#                             insfile='statistics.csv.ins', #optional, the instruction file name
#                             index_cols=["metric","name","constituent"], #column header to use as index; can also use column number (zero-based) instead of the header name
#                             use_cols=["value"], #names of columns that include observation values; can also use column
#                             prefix="metrics") #prefix to all observation names; choose something logical and easy to find. We use it later on to select observations
#%% Add post processing functions
#TODO: Figure out where the files should live prior to using this function. Ideally add functions from this script.
# pf.add_py_function("outputs.py",
#                    "mult_well_function(arg1='userarg')",
#                     is_pre_cmd = True)

#%% Build Pst Object
pst = pf.build_pst(template_ws /f'{model_name}.pst')
pst.add_parameters(template_ws / uci_tpl, template_ws / uci_file)
pst.parameter_data = parameter_data.reset_index()
pst.write(template_ws /f'{model_name}.pst')

control_file = f'{model_name}.pst'
#%% Run the model forward to generate outputs and test post processing functions
pst.write(os.path.join(template_ws, control_file),version=2)
# %%
