#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
Plot median aerosol size distribution at SGP
for cold and warm temperature conditions.

Cold condition = temp_degC <= 25th percentile
Warm condition = temp_degC >= 75th percentile

The middle 50% of the temperature distribution is excluded.

The ASD is reconstructed from the fitted three-mode lognormal parameters:
Nmode_1-3, Dpg_um_1-3, sigma_1-3.

Solid lines show the median reconstructed ASD.
Shading shows the 25th-75th percentile range.
"""

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


# =========================================================
# User settings
# =========================================================

path = "."
start_year = 2017
end_year = 2023
summer_months = [6, 7, 8]

site = "SGP"

asd_file = path + "/ARM_SGP_ASD_data/SPG_hourly_all_2017_2024_Apr_Oct.csv"
met_file = path + "/aerosol_forcing/SGP_ARM_met_2017_2023_pressure_temp_rh.csv"

output_folder = path + "/ASD_cold_warm_temperature_plots/"
os.makedirs(output_folder, exist_ok=True)


# =========================================================
# Functions
# =========================================================

def lognormal_mode_dNdlogDp(Dp_um, N, Dpg_um, sigma):
    """
    Reconstruct one lognormal aerosol mode.

    Returns dN/dlog10(Dp).

    Parameters
    ----------
    Dp_um : array
        Particle diameter grid in micrometers.
    N : float
        Total number concentration of the mode [cm^-3].
    Dpg_um : float
        Geometric mean diameter [um].
    sigma : float
        Geometric standard deviation.
    """

    Dp_um = np.asarray(Dp_um)

    return (
        N / (np.sqrt(2 * np.pi) * np.log10(sigma))
        * np.exp(
            - (np.log10(Dp_um / Dpg_um))**2
            / (2 * np.log10(sigma)**2)
        )
    )


def total_asd_distribution(row, Dp_um):
    """
    Sum the three fitted lognormal modes for one timestep.
    """

    total = np.zeros_like(Dp_um, dtype=float)

    for i in [1, 2, 3]:
        N = row[f"Nmode_{i}"]
        Dpg = row[f"Dpg_um_{i}"]
        sigma = row[f"sigma_{i}"]

        if (
            np.isfinite(N)
            and np.isfinite(Dpg)
            and np.isfinite(sigma)
            and N > 0
            and Dpg > 0
            and sigma > 1
        ):
            total += lognormal_mode_dNdlogDp(Dp_um, N, Dpg, sigma)

    return total


# =========================================================
# 1. Read SGP ASD data
# =========================================================

df_asd = pd.read_csv(asd_file, header=0, sep="\t")

df_asd = df_asd.rename(columns={"Time": "time"})
df_asd["time"] = pd.to_datetime(df_asd["time"], errors="coerce").dt.round("min")


# =========================================================
# 2. Clean fitted modal parameters
# =========================================================

df_asd[["sigma_1", "sigma_2", "sigma_3"]] = (
    df_asd[["sigma_1", "sigma_2", "sigma_3"]].replace(1.001, 1.100)
)

df_asd[["Dpg_um_1", "Dpg_um_2", "Dpg_um_3"]] = (
    df_asd[["Dpg_um_1", "Dpg_um_2", "Dpg_um_3"]].replace(0.001, 0.050)
)

for i in [1, 2, 3]:
    df_asd.loc[df_asd[f"Nmode_{i}"] <= 0.001, f"Dpg_um_{i}"] = 0.050

# Average duplicate timesteps, if present
df_asd = df_asd.groupby("time", as_index=False).mean(numeric_only=True)


# =========================================================
# 3. Read SGP ARM meteorology / temperature data
# =========================================================

# Your file has columns like:
# time, atmos_pressure_kPa, temp_degC, rh_percent
df_temp = pd.read_csv(met_file)

# Clean column names in case there are accidental spaces
df_temp.columns = df_temp.columns.str.strip()

df_temp["time"] = pd.to_datetime(df_temp["time"], errors="coerce")
df_temp["temp_degC"] = pd.to_numeric(df_temp["temp_degC"], errors="coerce")

# Keep only useful columns
df_temp = df_temp[["time", "temp_degC"]].copy()

# Remove bad values
df_temp = df_temp.dropna(subset=["time", "temp_degC"])

# Keep JJA 2017-2023
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()

# Hourly mean temperature, matching hourly ASD data
df_temp = (
    df_temp
    .set_index("time")
    .resample("H")
    .mean(numeric_only=True)
    .reset_index()
)

print("Temperature data points after filtering:", len(df_temp))
print(df_temp.head())


# =========================================================
# 4. Apply ASD quality and time filtering
# =========================================================

required_cols = [
    "N100",
    "r2",
    "Dpg_um_1", "Dpg_um_2", "Dpg_um_3",
    "sigma_1", "sigma_2", "sigma_3",
    "Nmode_1", "Nmode_2", "Nmode_3",
]

df_asd = df_asd.replace([np.inf, -np.inf], np.nan)
df_asd = df_asd.dropna(subset=required_cols)

df_asd = df_asd[
    (df_asd["r2"] >= 0.99)
    & (df_asd["N100"] > 0)
    & (df_asd[[
        "Dpg_um_1", "Dpg_um_2", "Dpg_um_3",
        "sigma_1", "sigma_2", "sigma_3",
        "Nmode_1", "Nmode_2", "Nmode_3"
    ]] > 0).all(axis=1)
].copy()

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

print("Total valid ASD points:", len(df_asd))


# =========================================================
# 5. Merge ASD and temperature
# =========================================================

df_asd["time"] = pd.to_datetime(df_asd["time"]).dt.floor("H")
df_temp["time"] = pd.to_datetime(df_temp["time"]).dt.floor("H")

df_merged = pd.merge(df_asd, df_temp, on="time", how="inner")

df_merged = df_merged.replace([np.inf, -np.inf], np.nan)
df_merged = df_merged.dropna(subset=required_cols + ["temp_degC"])

print("Collocated ASD-temperature points:", len(df_merged))


# =========================================================
# 6. Define cold and warm temperature conditions
# =========================================================

cold_temp_threshold = df_merged["temp_degC"].quantile(0.25)
warm_temp_threshold = df_merged["temp_degC"].quantile(0.75)

cold_data = df_merged[df_merged["temp_degC"] <= cold_temp_threshold].copy()
warm_data = df_merged[df_merged["temp_degC"] >= warm_temp_threshold].copy()

print("\nCold temperature condition:")
print(f"temp_degC <= 25th percentile = {cold_temp_threshold:.2f} °C")
print(f"N = {len(cold_data)}")
print(f"Mean temp_degC = {cold_data['temp_degC'].mean():.2f} °C")
print(f"Median temp_degC = {cold_data['temp_degC'].median():.2f} °C")
print(f"Mean N100 = {cold_data['N100'].mean():.1f} cm^-3")
print(f"Median N100 = {cold_data['N100'].median():.1f} cm^-3")

print("\nWarm temperature condition:")
print(f"temp_degC >= 75th percentile = {warm_temp_threshold:.2f} °C")
print(f"N = {len(warm_data)}")
print(f"Mean temp_degC = {warm_data['temp_degC'].mean():.2f} °C")
print(f"Median temp_degC = {warm_data['temp_degC'].median():.2f} °C")
print(f"Mean N100 = {warm_data['N100'].mean():.1f} cm^-3")
print(f"Median N100 = {warm_data['N100'].median():.1f} cm^-3")


# =========================================================
# 7. Reconstruct aerosol size distributions
# =========================================================

# Diameter grid: 3 nm to 1000 nm
Dp_um = np.logspace(np.log10(0.003), np.log10(1.0), 300)

cold_dist = np.array([
    total_asd_distribution(row, Dp_um)
    for _, row in cold_data.iterrows()
])

warm_dist = np.array([
    total_asd_distribution(row, Dp_um)
    for _, row in warm_data.iterrows()
])


# =========================================================
# 8. Median ASD and 25-75 percentile range
# =========================================================

cold_center = np.nanmedian(cold_dist, axis=0)
warm_center = np.nanmedian(warm_dist, axis=0)

cold_q25, cold_q75 = np.nanpercentile(cold_dist, [25, 75], axis=0)
warm_q25, warm_q75 = np.nanpercentile(warm_dist, [25, 75], axis=0)


# =========================================================
# 9. Plot cold vs warm ASD
# =========================================================

fig, ax = plt.subplots(figsize=(7.6, 5.4))

# Cold temperature condition
ax.plot(
    Dp_um * 1000,
    cold_center,
    color="royalblue",
    lw=2.8,
    label=fr"Cold: $T \leq$ {cold_temp_threshold:.1f}$^\circ$C"
)

ax.fill_between(
    Dp_um * 1000,
    cold_q25,
    cold_q75,
    color="royalblue",
    alpha=0.18,
    linewidth=0
)

# Warm temperature condition
ax.plot(
    Dp_um * 1000,
    warm_center,
    color="crimson",
    lw=2.8,
    label=fr"Warm: $T \geq$ {warm_temp_threshold:.1f}$^\circ$C"
)

ax.fill_between(
    Dp_um * 1000,
    warm_q25,
    warm_q75,
    color="crimson",
    alpha=0.18,
    linewidth=0
)


# =========================================================
# 10. Axis settings
# =========================================================

ax.set_xscale("log")

# Keep y-axis linear
ax.set_ylim(bottom=0)

# Give some headroom above the highest percentile shading
ymax = np.nanmax([
    np.nanmax(cold_q75),
    np.nanmax(warm_q75),
    np.nanmax(cold_center),
    np.nanmax(warm_center)
])

ax.set_ylim(0, ymax * 1.15)

ax.set_xlabel(r"Particle diameter, $D_p$ (nm)", fontsize=15)
ax.set_ylabel(r"dN/dlog$_{10}D_p$ (cm$^{-3}$)", fontsize=15)

ax.set_title(site, fontsize=18, pad=10)


# =========================================================
# 11. Style
# =========================================================

ax.grid(True, which="both", linestyle="--", alpha=0.25)

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

ax.tick_params(
    axis="both",
    which="minor",
    direction="in",
    length=3,
    width=1.0
)

ax.legend(
    frameon=True,
    fontsize=11,
    loc="upper right"
)

plt.tight_layout()

output_file = (
    output_folder
    + f"ASD_cold_vs_warm_temperature_{site}_{start_year}_{end_year}_JJA_25_75_median_IQR.png"
)

plt.savefig(output_file, dpi=300, bbox_inches="tight")
plt.show()

print("\nSaved figure:")
print(output_file)