from importlib import resources

import matplotlib.pyplot as plt
import numpy as np
import pandas

from pyrealm.pmodel import (
    PModelEnvironment,
    SubdailyScaler,
)
from pyrealm.pmodel.new_pmodel import PModelNew, SubdailyPModelNew

env = PModelEnvironment(
    tc=np.array([20]),
    vpd=np.array([1000]),
    co2=np.array([400]),
    patm=np.array([101325.0]),
)

p = PModelNew(env=env, fapar=np.array([1]), ppfd=np.array([300]))
print(p)
p.summarize()

data_path = resources.files("pyrealm_build_data.subdaily") / "subdaily_BE_Vie_2014.csv"

data = pandas.read_csv(str(data_path))

# Extract the key half hourly timestep variables as numpy arrays
temp_subdaily = data["ta"].to_numpy()
vpd_subdaily = data["vpd"].to_numpy()
co2_subdaily = data["co2"].to_numpy()
patm_subdaily = data["patm"].to_numpy()
ppfd_subdaily = data["ppfd"].to_numpy()
fapar_subdaily = data["fapar"].to_numpy()
datetime_subdaily = pandas.to_datetime(data["time"]).to_numpy()


subdaily_env = PModelEnvironment(
    tc=temp_subdaily,
    vpd=vpd_subdaily,
    co2=co2_subdaily,
    patm=patm_subdaily,
)

# Fit the standard P Model
pmodel_standard = PModelNew(
    subdaily_env,
    method_kphio="fixed",
    reference_kphio=1 / 8,
    ppfd=ppfd_subdaily,
    fapar=fapar_subdaily,
)
pmodel_standard._fit_model()
pmodel_standard.summarize()


# Create the fast slow scaler

fsscaler = SubdailyScaler(datetime_subdaily)

# Set the acclimation window as the values within a one hour window centred on noon
fsscaler.set_window(
    window_center=np.timedelta64(12, "h"),
    half_width=np.timedelta64(30, "m"),
)

# Fit the P Model with fast and slow responses
pmodel_subdaily = SubdailyPModelNew(
    env=subdaily_env,
    fs_scaler=fsscaler,
    allow_holdover=True,
    ppfd=ppfd_subdaily,
    fapar=fapar_subdaily,
    reference_kphio=1 / 8,
)

pmodel_subdaily._fit_model()
pmodel_subdaily.summarize()


idx = np.arange(48 * 120, 48 * 130)
plt.figure(figsize=(10, 4))
plt.plot(datetime_subdaily[idx], pmodel_standard.gpp[idx], label="Instantaneous model")
plt.plot(datetime_subdaily[idx], pmodel_subdaily.gpp[idx], "r-", label="Slow responses")
plt.ylabel("GPP")
plt.legend(frameon=False)
plt.show()