import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import statsmodels.api as stm
# import pymc3 as pm
import stan
import nest_asyncio
from scipy.stats import gaussian_kde, pearsonr
nest_asyncio.apply()


# Years and months
start_year = 2012
end_year = 2023
summer_months = [7, 8]

# =========================================================
# 1. READ AERONET DATA
# =========================================================
df_aer = pd.read_csv("aeronet_hyytiälä_AOD_AE_2012_2023_JJA_notSDA.csv")

df_aer["time"] = pd.to_datetime(
    df_aer["Date(dd:mm:yyyy)"] + " " + df_aer["Time(hh:mm:ss)"],
    errors="coerce"
)

# Keep only needed columns
df_aer = df_aer[[
    "time",
    "AOD_500nm",
    "440-870_Angstrom_Exponent"
]].copy()

# Filter: 2012–2023, July-August only
df_aer = df_aer[
    (df_aer["time"].dt.year >= start_year) &
    (df_aer["time"].dt.year <= end_year) &
    (df_aer["time"].dt.month.isin(summer_months))
].copy()

# Remove bad / cloudy points
df_aer = df_aer[
    np.isfinite(df_aer["AOD_500nm"]) &
    np.isfinite(df_aer["440-870_Angstrom_Exponent"]) &
    (df_aer["AOD_500nm"] > 0) &
    (df_aer["440-870_Angstrom_Exponent"] > 0.75)
].copy()

df_aer["date"] = df_aer["time"].dt.floor("D")
df_aer["hour"] = df_aer["time"].dt.hour
df_aer["year"] = df_aer["time"].dt.year

# =========================================================
# 2. DAILY REPRESENTATIVE AERONET MEANS
# =========================================================
daily_aod_rows = []

for date, g in df_aer.groupby("date"):
    hh = g["hour"].to_numpy()

    n_total = len(g)
#    n_47 = np.sum((hh >= 4) & (hh <= 7))
#    n_811 = np.sum((hh >= 8) & (hh <= 11))
#    n_1215 = np.sum((hh >= 12) & (hh <= 15))

    if (n_total >= 6): # and (n_47 >= 2) and (n_811 >= 2) and (n_1215 >= 2):
        daily_aod_rows.append({
            "date": date,
            "year": int(g["year"].iloc[0]),
            "AOD_500nm": np.nanmean(g["AOD_500nm"])
        })

daily_aod = pd.DataFrame(daily_aod_rows).sort_values("date").reset_index(drop=True)

# =========================================================
# 3. DAILY MEAN TEMPERATURE
# =========================================================
df_temp = pd.read_csv("smear2_hyytiala_air_temp_2012_2023.csv")

df_temp["time"] = pd.to_datetime(
    df_temp[["Year", "Month", "Day", "Hour", "Minute", "Second"]]
)

df_temp = df_temp.rename(columns={"HYY_META.T168": "T"})
df_temp = df_temp[["time", "T"]].copy()

df_temp = df_temp[
    (df_temp["time"].dt.year >= start_year) &
    (df_temp["time"].dt.year <= end_year) &
    (df_temp["time"].dt.month.isin(summer_months))
].copy()

df_temp["date"] = df_temp["time"].dt.floor("D")
df_temp["year"] = df_temp["time"].dt.year

daily_temp = (
    df_temp.groupby("date", as_index=False)
    .agg({"T": "mean", "year": "first"})
    .rename(columns={"T": "temperature_C"})
)

# =========================================================
# 4. COLLOCATE DAILY TEMPERATURE AND DAILY AOD
# =========================================================
df = pd.merge(daily_aod, daily_temp, on=["date", "year"], how="inner")
df = df.dropna(subset=["AOD_500nm", "temperature_C"]).copy()

# Keep positive temperatures to match old model assumption trueT>0
df = df[df["temperature_C"] > 0].copy()

# =========================================================
# 5. SUMMER MEANS BY YEAR
# =========================================================
summer = (
    df.groupby("year", as_index=False)
    .agg({
        "temperature_C": "mean",
        "AOD_500nm": "mean"
    })
)

# Define x and y for Bayesian and OLS fits using daily means for all years
x = df["temperature_C"].to_numpy()
y = df["AOD_500nm"].to_numpy()


## =========================================================
## 9. Stable Stan model
## =========================================================
#stan_code = """
#data {
#  int<lower=0> N;
#  vector[N] AODobs;
#  vector[N] Tobs;

#  real trueTmu;
#  real trueTsigma;

#  real meanA;
#  real stdA;
#  real meanB;
#  real stdB;
#}
#parameters {
#  real a;
#  real b;
#  vector<lower=0.0>[N] trueT;

#  real<lower=1e-6> sigma_AOD;
#  real<lower=1e-6> sigma_T;
#}
#model {
#  trueT ~ normal(trueTmu, trueTsigma);

#  AODobs ~ normal(a * trueT + b, sigma_AOD);
#  Tobs   ~ normal(trueT, sigma_T);

#  a ~ normal(meanA, stdA);
#  b ~ normal(meanB, stdB);

#  sigma_AOD ~ normal(0.05, 0.05);
#  sigma_T   ~ normal(2.0, 1.0);
#}
#"""

#stan_data = {
#    "N": len(x),
#    "AODobs": y,
#    "Tobs": x,
#    "trueTmu": float(np.mean(x)),
#    "trueTsigma": float(np.std(x)),
#    "meanA": 0.0,
#    "stdA": 10.0,
#    "meanB": 0.0,
#    "stdB": 10.0,
#}

#posterior = stan.build(stan_code, data=stan_data, random_seed=1)
#fit = posterior.sample(num_chains=4, num_samples=1000)

## Extract posterior samples
#a_mcmc = fit["a"].reshape(-1)
#b_mcmc = fit["b"].reshape(-1)

## Compute Bayesian slope and intercept
#bayes_slope = np.mean(a_mcmc)
#bayes_intercept = np.mean(b_mcmc)

#print("Bayesian slope:", bayes_slope)
#print("Bayesian slope 95% CrI:", np.percentile(a_mcmc, [2.5, 97.5]))

## Compute credible intervals for Bayesian fit
#x_fit = np.linspace(np.min(x), np.max(x), 250)
#bayes_lower, bayes_upper = get_credible_interval_limits(a_mcmc, b_mcmc, x_fit, p=95.0)

## OLS fit for daily medians
#X = stm.add_constant(x)
#ols_fit = stm.OLS(y, X).fit()

#ols_intercept = ols_fit.params[0]
#ols_slope = ols_fit.params[1]

## Extract OLS slope confidence intervals
#ols_conf_int = ols_fit.conf_int()
#ols_slope_ci_lower = ols_conf_int[1, 0]  # Lower bound for slope
#ols_slope_ci_upper = ols_conf_int[1, 1]  # Upper bound for slope

## Compute OLS fit line
#y_fit_ols = ols_slope * x_fit + ols_intercept
#pred = ols_fit.get_prediction(stm.add_constant(x_fit))
#pred_ci = pred.summary_frame(alpha=0.05)

## =========================================================
## Helper Function for Credible Intervals
## =========================================================
#def get_credible_interval_limits(a_mcmc, b_mcmc, xx, p=95.0):
#    datapoints = (
#        a_mcmc.ravel()[np.newaxis, :] * xx.ravel()[:, np.newaxis]
#        + b_mcmc.ravel()[np.newaxis, :]
#    )
#    lower = np.percentile(datapoints, (100.0 - p) / 2.0, axis=1)
#    upper = np.percentile(datapoints, 100.0 - (100.0 - p) / 2.0, axis=1)
#    return lower, upper

## =========================================================
## 11. Plot
## =========================================================
#plt.figure(figsize=(8, 6))

## Daily scatter points
#plt.plot(
#    df["temperature_C"].values,
#    df["AOD_500nm"].values,
#    "o",
#    color="lightgray",
#    alpha=0.6,
#    label="Daily means"
#)

## Summer means as squares
#plt.plot(
#    summer["temperature_C"].values,
#    summer["AOD_500nm"].values,
#    marker="s",
#    mew=1,
#    mec="k",
#    markerfacecolor="turquoise",
#    markersize=11,
#    ls="none",
#    label="Summer means"
#)

## OLS
#plt.plot(
#    x_fit, y_fit_ols,
#    color="blue",
#    lw=2,
#    ls="--",
#    label=f"OLS fit y={ols_slope:.3f}x{ols_intercept:+.3f}"
#)
#plt.fill_between(
#    x_fit,
#    pred_ci["mean_ci_lower"].values,
#    pred_ci["mean_ci_upper"].values,
#    color="lightblue",
#    alpha=0.15
#)

## Bayesian
#y_fit_bayes = bayes_slope * x_fit + bayes_intercept
#plt.plot(
#    x_fit, y_fit_bayes,
#    color="limegreen",
#    lw=2,
#    ls="--",
#    label=f"Bayesian fit y={bayes_slope:.3f}x{bayes_intercept:+.3f}"
#)
#plt.fill_between(
#    x_fit,
#    bayes_lower,
#    bayes_upper,
#    facecolor="palegreen",
#    alpha=0.4,
#    interpolate=True
#)

#plt.title("Hyytiälä", fontsize=18)
## plt.xlim(5, 30)
#plt.gca().xaxis.set_major_formatter(plt.FuncFormatter(lambda x, _: f"{x:.0f}"))
#plt.tick_params(axis="both", which="major", direction="in", length=6, width=2, labelsize=14)
#plt.tick_params(axis="both", which="minor", labelsize=14)

#plt.ylabel(r"AOD at 500 nm", fontsize=16)
#plt.xlabel(r"Temperature ($^\circ$C)", fontsize=16)

#leg = plt.legend(frameon=True, fontsize=12, loc=2)
#frame = leg.get_frame()
#frame.set_facecolor("w")
#frame.set_edgecolor("k")

#plt.tight_layout()
#plt.savefig("Hyytiala_T_vs_AERONET_total_AOD500_daily_and_summer_means_2012_2023.png", dpi=150)
#plt.show()


# =========================================================
# 6. AOD - Temperature point-density plot
# =========================================================

# ---------------------------------------------------------
# Prepare daily data
# ---------------------------------------------------------
df_plot = df.copy()

df_plot = df_plot[
    np.isfinite(df_plot["temperature_C"]) &
    np.isfinite(df_plot["AOD_500nm"]) &
    (df_plot["AOD_500nm"] > 0)
].copy()

x = df_plot["temperature_C"].to_numpy()
y = df_plot["AOD_500nm"].to_numpy()

print("Number of valid daily AOD-temperature points:", len(df_plot))

# ---------------------------------------------------------
# OLS regression: AOD = intercept + slope * T
# ---------------------------------------------------------
X = stm.add_constant(x)
ols_fit = stm.OLS(y, X).fit()

ols_intercept = ols_fit.params[0]
ols_slope = ols_fit.params[1]

ols_conf_int = ols_fit.conf_int(alpha=0.05)
ols_slope_ci_lower = ols_conf_int[1, 0]
ols_slope_ci_upper = ols_conf_int[1, 1]

# Smooth x-grid for regression line
x_fit = np.linspace(np.nanmin(x), np.nanmax(x), 250)
X_fit = stm.add_constant(x_fit)

pred = ols_fit.get_prediction(X_fit)
pred_mean = pred.predicted_mean
pred_ci = pred.conf_int(alpha=0.05)

# Pearson correlation
r_value, p_value = pearsonr(x, y)

print(f"OLS slope dAOD/dT = {ols_slope:.5f} °C^-1")
print(f"95% CI = [{ols_slope_ci_lower:.5f}, {ols_slope_ci_upper:.5f}] °C^-1")
print(f"r = {r_value:.2f}, p = {p_value:.3g}")

# ---------------------------------------------------------
# Point density for daily points
# ---------------------------------------------------------
xy = np.vstack([x, y])
density = gaussian_kde(xy)(xy)

# Plot low-density points first, high-density points last
idx = density.argsort()
x_sorted = x[idx]
y_sorted = y[idx]
density_sorted = density[idx]

# ---------------------------------------------------------
# Prepare summer means
# ---------------------------------------------------------
summer_plot = summer.copy()

summer_plot = summer_plot[
    np.isfinite(summer_plot["temperature_C"]) &
    np.isfinite(summer_plot["AOD_500nm"])
].copy()

# =========================================================
# 7. Plot
# =========================================================

fig, ax = plt.subplots(figsize=(9.2, 6.4))

# ---------------------------------------------------------
# Daily means coloured by point density
# ---------------------------------------------------------
sc = ax.scatter(
    x_sorted,
    y_sorted,
    c=density_sorted,
    cmap="RdPu",
    s=58,
    alpha=0.78,
    edgecolor="none"
)

cbar = fig.colorbar(sc, ax=ax, pad=0.02)
cbar.set_label("Point density", fontsize=13)
cbar.ax.tick_params(labelsize=11)

# ---------------------------------------------------------
# 95% CI curtain
# ---------------------------------------------------------
ci_band = ax.fill_between(
    x_fit,
    pred_ci[:, 0],
    pred_ci[:, 1],
    color="red",
    alpha=0.18,
    label="95% CI of fit"
)

# ---------------------------------------------------------
# OLS regression line
# ---------------------------------------------------------
ols_line, = ax.plot(
    x_fit,
    pred_mean,
    color="red",
    lw=2.8,
    ls="--",
    label=(
        r"$\mathrm{AOD}_{500} = "
        f"{ols_slope:.2e}T "
        f"{ols_intercept:+.3f}$"
        )
)

# ---------------------------------------------------------
# Summer means as dark-purple squares, no error bars
# ---------------------------------------------------------
summer_handle = ax.scatter(
    summer_plot["temperature_C"],
    summer_plot["AOD_500nm"],
    s=130,
    marker="s",
    facecolor="#4B0082",
    edgecolor="black",
    linewidth=1.1,
    alpha=0.95,
    zorder=7,
    label="Summer means"
)

# ---------------------------------------------------------
# Labels and style
# ---------------------------------------------------------
ax.set_xlabel(r"Temperature ($^\circ$C)", fontsize=16)
ax.set_ylabel(r"AOD at 500 nm", fontsize=16)

ax.set_title(
    r"Hyytiälä",
    fontsize=17,
    pad=12
)

ax.tick_params(
    axis="both",
    which="major",
    direction="in",
    length=6,
    width=1.8,
    labelsize=13
)

ax.grid(True, linestyle="--", alpha=0.22)

ax.xaxis.set_major_formatter(
    plt.FuncFormatter(lambda val, _: f"{val:.0f}")
)

# ---------------------------------------------------------
# One legend box only
# ---------------------------------------------------------
leg = ax.legend(
    handles=[ols_line, ci_band, summer_handle],
    frameon=True,
    fontsize=10.5,
    loc="upper left",
    handlelength=2.5,
    borderpad=0.8,
    labelspacing=0.8
)

leg.get_frame().set_facecolor("white")
leg.get_frame().set_edgecolor("black")
leg.get_frame().set_alpha(0.95)

plt.tight_layout()

plt.savefig(
    "Hyytiala_AOD500_vs_T_density_daily_and_summer_means_2012_2023_nTotal_only.png",
    dpi=300,
    bbox_inches="tight"
)

plt.show()


## =========================
## 5. DIRECT RADIATIVE EFFECT / FORCING
## =========================
## Formula from Yli-Juuti et al. setup
## DRF = Srad * phi * tau_DRE_clear_Wm2a * (1 - Cc) * Tatm^2 * (1 - Rs)^2 *
##       ( 2*Rs*(1-w)/((1-Rs)^2) - beta*w )

Srad = 398.0
phi = 1.82
Tatm = 0.76
Rs = 0.12
omega = 0.92
beta = 0.29

# Clear-sky and all-sky cloud fraction choices
Cc_clear = 0.0
Cc_allsky = 0.511  ## CF calculated from MODIS L2 data with fraction of water phase clouds in the (water_cloudy + clear_sky)

def drf_from_tau(tau_a, Cc):
    return (
        Srad * phi * tau_a * (1 - Cc) * (Tatm ** 2) * ((1 - Rs) ** 2) *
        (2 * Rs * (1 - omega) / ((1 - Rs) ** 2) - beta * omega)
    )

df["DRE_clear_Wm2"] = drf_from_tau(df["AOD_500nm"], Cc_clear)
df["DRE_allsky_Wm2"] = drf_from_tau(df["AOD_500nm"], Cc_allsky)

# Calculate mean DRF values for clear-sky and all-sky
mean_DRF_clear = df["DRE_clear_Wm2"].mean()
mean_DRF_allsky = df["DRE_allsky_Wm2"].mean()

# Create a summary DataFrame
summary_df = pd.DataFrame({
    "Mean_DRF_clear_Wm2": [mean_DRF_clear],
    "Mean_DRF_allsky_Wm2": [mean_DRF_allsky]
})

# Save the summary to CSV
summary_df.to_csv("mean_drf_summary.csv", index=False)
print("Saved: mean_drf_summary.csv")


# =========================
# 6. TEMPERATURE-DEPENDENT FEEDBACK STRENGTH
# =========================
# Simple linear slope d(tau_a)/dT
# You can replace this with statsmodels OLS if you want confidence intervals
x = df["temperature_C"].values
y = df["AOD_500nm"].values

mask = np.isfinite(x) & np.isfinite(y)
slope_tau_per_C, intercept = np.polyfit(x[mask], y[mask], 1)

# Radiative efficiency factor (W m^-2 per unit tau_a)
def radiative_efficiency(Cc):
    return (
        Srad * phi * (1 - Cc) * (Tatm ** 2) * ((1 - Rs) ** 2) *
        (2 * Rs * (1 - omega) / ((1 - Rs) ** 2) - beta * omega)
    )

RE_clear = radiative_efficiency(Cc_clear)
RE_allsky = radiative_efficiency(Cc_allsky)

# Temperature-dependent direct radiative feedback (W m^-2 C^-1)
DRF_dT_clear = RE_clear * slope_tau_per_C
DRF_dT_allsky = RE_allsky * slope_tau_per_C

# Extract OLS slope confidence intervals
ols_slope_ci_lower = ols_conf_int[1, 0]  # Lower bound for slope
ols_slope_ci_upper = ols_conf_int[1, 1]  # Upper bound for slope

# Compute OLS fit line
y_fit_ols = ols_slope * x_fit + ols_intercept
pred = ols_fit.get_prediction(stm.add_constant(x_fit))
pred_ci = pred.summary_frame(alpha=0.05)

# Calculate DRF confidence intervals for clear-sky and all-sky
DRF_dT_clear_ci_lower = RE_clear * ols_slope_ci_lower
DRF_dT_clear_ci_upper = RE_clear * ols_slope_ci_upper

DRF_dT_allsky_ci_lower = RE_allsky * ols_slope_ci_lower
DRF_dT_allsky_ci_upper = RE_allsky * ols_slope_ci_upper

# Print the results
print("Number of valid daily means:", len(df))
print("d(tau_a)/dT =", slope_tau_per_C, "per C")
print("Direct radiative feedback (clear-sky) [W m^-2 C^-1]:", DRF_dT_clear)
print("Direct radiative feedback (all-sky)  [W m^-2 C^-1]:", DRF_dT_allsky)
print("Direct radiative feedback (clear-sky) CI [W m^-2 C^-1]:",
      DRF_dT_clear_ci_lower, "to", DRF_dT_clear_ci_upper)
print("Direct radiative feedback (all-sky) CI [W m^-2 C^-1]:",
      DRF_dT_allsky_ci_lower, "to", DRF_dT_allsky_ci_upper)

# Save the results to a .txt file
with open("drf_results_summary.txt", "w") as f:
    f.write(f"Number of valid daily means: {len(df)}\n")
    f.write(f"d(tau_a)/dT = {slope_tau_per_C} per C\n")
    f.write(f"Direct radiative feedback (clear-sky) [W m^-2 C^-1]: {DRF_dT_clear}\n")
    f.write(f"Direct radiative feedback (all-sky)  [W m^-2 C^-1]: {DRF_dT_allsky}\n")
    f.write(f"Direct radiative feedback (clear-sky) CI [W m^-2 C^-1]: {DRF_dT_clear_ci_lower} to {DRF_dT_clear_ci_upper}\n")
    f.write(f"Direct radiative feedback (all-sky) CI [W m^-2 C^-1]: {DRF_dT_allsky_ci_lower} to {DRF_dT_allsky_ci_upper}\n")

print("Saved: drf_results_summary.txt")


# =========================
# 7. SAVE OUTPUT
# =========================
df.to_csv("daily_aeronet_temp_DRE.csv", index=False)
print("Saved: daily_aeronet_temp_DRE.csv")