#%%
from pathlib import Path
import pandas as pd
import numpy as np

# pyhcal imports
from pyhcal import calibrators
from pyhcal import PestWorkflow
from pyhcal.pest.observation_builder import ObservationBuilder
import pyemu
from pyhcal.pest.pest_workflow import PestWorkflow
#%%
project_path = 'C:\\Users\\mfratki\\Documents\\Projects\\Tests\\Nemadji'
cal = calibrators.calibrator(project_path)
cal.load_model(0)

#%%
project_path = 'C:\\Users\\mfratki\\Documents\\Projects\\Tests\\Nemadji'
pest = PestWorkflow(project_path)
pest.load_model(0)

pest.add_hspf_parameter('PERLND','PWAT-PARM2',0,'LZSN')
pest.add_hspf_parameter('PERLND','PWAT-PARM2',0,'INFILT')

name = 'outlet'
station_ids = ['E05011002']
reach_ids = [103]
constituent = 'Q'
time_step = 'D'
weight = 1.0

df = pest.cal.compare_simulated_observed(station_ids, reach_ids, constituent, time_step).dropna(subset=['observed'])
pest.add_observation(name,constituent,'timeseries',weight,df['observed'],df['simulated'])



df = pest.cal.compare_simulated_observed(station_ids, reach_ids, constituent, time_step).dropna(subset=['observed'])
ob = ObservationBuilder()
ob.add_location(name,df['observed'],df['simulated'],weight,constituent)
observations = ob.build()


#%%
cal = calibrators.calibrator(project_path)
cal.load_model(0)
#%% Construct parameter template file
block = 'PERLND'
table = 'PWAT-PARM2'
table_id = 0
parameter = 'LZSN'

uci = cal.uci
uci.add_parameter_template(block, table, table_id, parameter,single_template=False)
uci.write_tpl(new_tpl_path = cal.model_path / cal.uci.filepath.name.replace('.uci', '.tpl'))

#%% Construct observation instruction file
name = 'outlet'
station_ids = ['E05011002']
reach_ids = [103]
constituent = 'Q'
time_step = 'D'
weight = 1.0

df = pest.cal.compare_simulated_observed(station_ids, reach_ids, constituent, time_step).dropna(subset=['observed'])
ob = ObservationBuilder()
ob.add_location(name,df['observed'],df['simulated'],weight,constituent)
observations = ob.build()


#%%


obs_dict = observations[['obsnme','obsval']].set_index('obsnme').to_dict()['obsval']
create_instruction_file(observations, cal.model_path / "output.ins")
write_simulated_output(observations, cal.model_path / "output.txt")

#%% Build Pst Object
pst = pyemu.Pst.from_io_files(
    tpl_files=[ cal.model_path /"Nemadji_0.tpl"], in_files=[ cal.model_path /"Nemadji_0.uci"],
    ins_files=[ cal.model_path /"output.ins"], out_files=[ cal.model_path /"output.txt"]
)
pst.adjust_weights(observations[['obsnme','weight']].set_index('obsnme').to_dict()['weight'])

#%%
pest.add_hspf_parameter(block, table, table_id, parameter,single_template=False)

pyemu.pst_utils.generic_pst()

#%% Add observation data
name = 'outlet'
station_ids = ['E05011002']
reach_ids = [103]
constituent = 'Q'
time_step = 'D'
df = pest.cal.compare_simulated_observed(station_ids, reach_ids, constituent, time_step).dropna(subset=['observed'])
pest.add_observation(name, constituent, observed=df['observed'], simulated=df['simulated'])

#%%
pest.build()
pest.write("my_case.pst")


#%%

from pyhcal.pest import PestWorkflow, create_instruction_file, write_simulated_output
#%%
project_path = 'C:\\Users\\mfratki\\Documents\\Projects\\Tests\\Nemadji'

pest = PestWorkflow(project_path)
pest.load_model(0)

# 1. Create template via UCI
cal.uci.add_parameter_template("PERLND", "PWAT-PARM2", 0, "LZSN", single_template=False)
cal.uci.write_tpl("model.tpl")

# 2. Build observation data
from pyhcal.pest import build_observations
obs_data = build_observations(observed, simulated, prefix="st1")

# 3. Create instruction file + initial output
create_instruction_file(obs_data, "output.ins")
write_simulated_output(obs_data, "output.txt")

# 4. Use pyemu to build PST
import pyemu
pst = pyemu.Pst.from_io_files(
    tpl_files=["model.tpl"], in_files=["model.uci"],
    ins_files=["output.ins"], out_files=["output.txt"]
)