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

"""
Density plot between AERONET AOD500 and ln(N100)
at Hyytiälä.

AERONET and ASD/N100 are collocated using nearest-time matching
within ±30 minutes.

Period: July-August, 2012-2018

Screening:
- AOD500 > 0
- Angstrom exponent 440-870 nm > 0.75
- ASD r2 >= 0.99
- Nt < 8000
- N100 > 0

Active plot:
- AOD500 vs ln(N100)

Commented alternative:
- AOD500 vs N100
"""

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from scipy.stats import pearsonr, gaussian_kde


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

aeronet_file = "aeronet_hyytiälä_AOD_AE_2012_2023_JJA_notSDA.csv"
asd_file = "Hyytiala_hourly_2012_2023_clean_ASD.csv"

start_year = 2012
end_year = 2018
summer_months = [7, 8]

output_file = (
    "AOD500_vs_lnN100_Hyytiala_AERONET_ASD_"
    "2012_2018_JulAug_density_viridis.png"
)


# =========================================================
# 1. Read AERONET data
# =========================================================

df_aero = pd.read_csv(aeronet_file)

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

df_aero = df_aero[[
    "time",
    "AOD_500nm",
    "440-870_Angstrom_Exponent"
]].copy()

df_aero = df_aero.dropna(
    subset=[
        "time",
        "AOD_500nm",
        "440-870_Angstrom_Exponent"
    ]
)

# Keep July-August, 2012-2018
df_aero = df_aero[
    (df_aero["time"].dt.year >= start_year) &
    (df_aero["time"].dt.year <= end_year) &
    (df_aero["time"].dt.month.isin(summer_months))
].copy()

# AERONET screening
df_aero = df_aero[
    (df_aero["AOD_500nm"] > 0) &
    (df_aero["440-870_Angstrom_Exponent"] > 0.75)
].copy()

df_aero = df_aero.sort_values("time").reset_index(drop=True)

print("AERONET points after filtering:", len(df_aero))


# =========================================================
# 2. Read ASD / N100 data
# =========================================================

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

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

# 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
df_asd = df_asd.groupby("time", as_index=False).mean(numeric_only=True)

# ASD quality filtering
valid_indices = np.where(
    (df_asd["r2"] >= 0.99) &
    (df_asd["Nt"] < 8000) &
    (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)
)[0]

df_asd = df_asd.iloc[valid_indices].copy()

# Keep July-August, 2012-2018
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()

df_asd = df_asd.dropna(subset=["time", "N100"])
df_asd = df_asd[df_asd["N100"] > 0].copy()

df_asd = df_asd.sort_values("time").reset_index(drop=True)

print("ASD/N100 points after filtering:", len(df_asd))


# =========================================================
# 3. Collocate AERONET with N100 within ±30 min
# =========================================================

colloc = pd.merge_asof(
    df_aero,
    df_asd,
    on="time",
    direction="nearest",
    tolerance=pd.Timedelta("30min")
)

colloc = colloc.dropna(subset=["N100", "AOD_500nm"]).copy()

print("Collocated AERONET-ASD points:", len(colloc))


# =========================================================
# 4. Prepare data
# =========================================================

plot_df = colloc[
    np.isfinite(colloc["AOD_500nm"]) &
    np.isfinite(colloc["N100"]) &
    (colloc["AOD_500nm"] > 0) &
    (colloc["N100"] > 0)
].copy()

# =========================================================
# ACTIVE VERSION: use ln(N100)
# =========================================================

plot_df["ln_N100"] = np.log(plot_df["N100"])

plot_df = plot_df[
    np.isfinite(plot_df["ln_N100"]) &
    np.isfinite(plot_df["AOD_500nm"])
].copy()

x = plot_df["ln_N100"].values
y = plot_df["AOD_500nm"].values

x_label = r"$\ln(N_{100})$ [cm$^{-3}$]"
output_file = (
    "AOD500_vs_lnN100_Hyytiala_AERONET_ASD_"
    "2012_2018_JulAug_density_viridis.png"
)


# =========================================================
# COMMENTED ALTERNATIVE: use N100 directly
# To use this, comment the ln(N100) block above and uncomment below.
# =========================================================

# plot_df = plot_df[
#     np.isfinite(plot_df["N100"]) &
#     np.isfinite(plot_df["AOD_500nm"])
# ].copy()
#
# x = plot_df["N100"].values
# y = plot_df["AOD_500nm"].values
#
# x_label = r"$N_{100}$ (cm$^{-3}$)"
# output_file = (
#     "AOD500_vs_N100_Hyytiala_AERONET_ASD_"
#     "2012_2018_JulAug_density_viridis.png"
# )


print("Final points used in density plot:", len(plot_df))

if len(plot_df) < 3:
    raise ValueError("Too few collocated points for correlation analysis.")


# =========================================================
# 5. Pearson correlation
# =========================================================

pearson_r, pearson_p = pearsonr(x, y)

print("\nCorrelation result:")
print(f"Pearson r = {pearson_r:.3f}, p = {pearson_p:.3g}")


# =========================================================
# 6. Density calculation
# =========================================================

xy = np.vstack([x, y])
density = gaussian_kde(xy)(xy)

# Sort points so high-density points are plotted on top
idx = density.argsort()
x_sorted = x[idx]
y_sorted = y[idx]
density_sorted = density[idx]


# =========================================================
# 7. Density-colored scatter plot
# =========================================================

fig, ax = plt.subplots(figsize=(7.4, 5.6))

sc = ax.scatter(
    x_sorted,
    y_sorted,
    c=density_sorted,
    s=18,
    cmap="viridis",
    alpha=0.85,
    edgecolor="none"
)

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


# =========================================================
# 8. Axis limits: include all points with padding
# =========================================================

x_min = np.nanmin(x)
x_max = np.nanmax(x)
y_min = np.nanmin(y)
y_max = np.nanmax(y)

x_pad = 0.06 * (x_max - x_min)
y_pad = 0.10 * (y_max - y_min)

ax.set_xlim(x_min - x_pad, x_max + x_pad)
ax.set_ylim(max(0, y_min - y_pad), y_max + y_pad)


# =========================================================
# 9. Pearson correlation in legend only
# =========================================================

ax.plot(
    [],
    [],
    color="none",
    label=fr"Pearson $r$ = {pearson_r:.2f}"
)


# =========================================================
# 10. Labels and style
# =========================================================

ax.set_xlabel(x_label, fontsize=15)
ax.set_ylabel(r"AOD$_{500}$", fontsize=15)

ax.set_title("Hyytiälä", fontsize=18, pad=10)

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

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

ax.legend(
    frameon=True,
    fontsize=12,
    loc="upper left"
)

plt.tight_layout()

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

plt.show()

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