import os
import numpy as np
import pandas as pd
import statsmodels.api as sm
import matplotlib.pyplot as plt

# ======================================================
# USER SETTINGS
# ======================================================

# NO2 CSV 
no2_csv = "/home/taneja/Documents/kanika/work/NEW_2025/SGP_new_local_filter/aerosol_forcing/OMI_SGP_NO2_daily_JJA_2017_2023.csv"
#no2_csv = "/media/taneja/Expansion/backup_home_May2026/Documents/kanika/work/NEW_2025/SGP_new_local_filter/aerosol_forcing/OMI_SGP_NO2_daily_JJA_2017_2023.csv"

# Your uploaded/input files
aeronet_csv = "/home/taneja/Documents/kanika/work/NEW_2025/SGP_new_local_filter/aerosol_forcing/aeronet_sgp_AOD_AE_2017_2023_JJA_notSDA.csv"
#aeronet_csv = "/media/taneja/Expansion/backup_home_May2026/Documents/kanika//work/NEW_2025/SGP_new_local_filter/aerosol_forcing/aeronet_sgp_AOD_AE_2017_2023_JJA_notSDA.csv"
arm_temp_csv = "/home/taneja/Documents/kanika/work/NEW_2025/SGP_new_local_filter/aerosol_forcing/SGP_ARM_met_2017_2023_pressure_temp_rh.csv"
#arm_temp_csv = "/media/taneja/Expansion/backup_home_May2026/Documents/kanika//work/NEW_2025/SGP_new_local_filter/aerosol_forcing/SGP_ARM_met_2017_2023_pressure_temp_rh.csv"
# Output directory
#out_dir = "/home/taneja/Documents/kanika/work/NEW_2025/SGP_new_local_filter/aerosol_forcing/"
out_dir = "/media/taneja/Expansion/backup_home_May2026/Documents/kanika/work/NEW_2025/SGP_new_local_filter/aerosol_forcing_TM/"
os.makedirs(out_dir, exist_ok=True)

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

# ======================================================
# 1. READ NO2 DAILY CSV
# ======================================================

df_no2 = pd.read_csv(no2_csv)
df_no2["date"] = pd.to_datetime(df_no2["date"]).dt.floor("D")

# Keep only required columns
df_no2 = df_no2[["date", "NO2_mean"]].copy()
df_no2 = df_no2.rename(columns={"NO2_mean": "NO2"})

# Keep JJA 2017-2023
df_no2 = df_no2[
    (df_no2["date"].dt.year >= start_year) &
    (df_no2["date"].dt.year <= end_year) &
    (df_no2["date"].dt.month.isin(summer_months))
].copy()

df_no2 = df_no2.dropna(subset=["NO2"])

print("NO2 daily data:")
print(df_no2.head())
print(f"NO2 rows: {len(df_no2)}")

# ======================================================
# 2. READ AERONET AOD AND MAKE DAILY MEANS
# ======================================================

df_aod = pd.read_csv(aeronet_csv)

# Combine AERONET date and time
df_aod["time"] = pd.to_datetime(
    df_aod["Date(dd:mm:yyyy)"] + " " + df_aod["Time(hh:mm:ss)"],
    format="%Y-%m-%d %H:%M:%S"
)

# Convert to daily mean AOD
df_aod["date"] = df_aod["time"].dt.floor("D")
df_aod_daily = (
    df_aod.groupby("date", as_index=False)
    .agg(
        AOD_500nm=("AOD_500nm", "mean"),
        AE_440_870=("440-870_Angstrom_Exponent", "mean"),
        n_aod_obs=("AOD_500nm", "count")
    )
)

# Keep JJA 2017-2023
df_aod_daily = df_aod_daily[
    (df_aod_daily["date"].dt.year >= start_year) &
    (df_aod_daily["date"].dt.year <= end_year) &
    (df_aod_daily["date"].dt.month.isin(summer_months))
].copy()

print("\nAERONET daily data:")
print(df_aod_daily.head())
print(f"AOD daily rows: {len(df_aod_daily)}")

# ======================================================
# 3. READ ARM TEMPERATURE AND MAKE DAILY MEANS
# ======================================================

df_temp = pd.read_csv(arm_temp_csv)
df_temp["time"] = pd.to_datetime(df_temp["time"])

# Convert to daily mean temperature
df_temp["date"] = df_temp["time"].dt.floor("D")
df_temp_daily = (
    df_temp.groupby("date", as_index=False)
    .agg(
        temp_degC=("temp_degC", "mean"),
        atmos_pressure_kPa=("atmos_pressure_kPa", "mean"),
        rh_percent=("rh_percent", "mean"),
        n_temp_obs=("temp_degC", "count")
    )
)

# Keep JJA 2017-2023
df_temp_daily = df_temp_daily[
    (df_temp_daily["date"].dt.year >= start_year) &
    (df_temp_daily["date"].dt.year <= end_year) &
    (df_temp_daily["date"].dt.month.isin(summer_months))
].copy()

print("\nARM daily temperature data:")
print(df_temp_daily.head())
print(f"Temperature daily rows: {len(df_temp_daily)}")

# ======================================================
# 4. COLLOCATE ALL THREE DATASETS BY DATE
# ======================================================

df_merged = df_no2.merge(df_aod_daily, on="date", how="inner")
df_merged = df_merged.merge(df_temp_daily, on="date", how="inner")

# Drop missing values in variables used for regression
df_merged = df_merged.dropna(subset=["NO2", "AOD_500nm", "temp_degC"]).copy()
df_merged = df_merged.sort_values("date").reset_index(drop=True)

print("\nMerged collocated dataset:")
print(df_merged.head())
print(f"Merged rows: {len(df_merged)}")

merged_csv = os.path.join(out_dir, "NO2_AOD_Temp_daily_JJA_2017_2023_merged.csv")
df_merged.to_csv(merged_csv, index=False)
print(f"Saved merged CSV: {merged_csv}")

# ======================================================
# 5. MULTIVARIATE FIT
#    AOD = a*NO2 + b*Temp + d
# ======================================================

X = df_merged[["NO2", "temp_degC"]].copy()
X = sm.add_constant(X)
y = df_merged["AOD_500nm"]

model = sm.OLS(y, X, missing="drop").fit()

print("\nMultivariate regression summary:")
print(model.summary())

# Save coefficients
coef_df = pd.DataFrame({
    "term": model.params.index,
    "coefficient": model.params.values,
    "std_error": model.bse.values,
    "t_value": model.tvalues.values,
    "p_value": model.pvalues.values
})

coef_csv = os.path.join(out_dir, "multivariate_fit_AOD_NO2_Temp_coefficients.csv")
coef_df.to_csv(coef_csv, index=False)
print(f"Saved coefficients CSV: {coef_csv}")

# Add fitted values and residuals
df_merged["AOD_fitted"] = model.predict(X)
df_merged["residual"] = df_merged["AOD_500nm"] - df_merged["AOD_fitted"]

fit_csv = os.path.join(out_dir, "multivariate_fit_AOD_NO2_Temp_fitted_timeseries.csv")
df_merged.to_csv(fit_csv, index=False)
print(f"Saved fitted data CSV: {fit_csv}")

# ======================================================
# 6. OBSERVED VS FITTED PLOT
# ======================================================

plt.figure(figsize=(6, 6))
plt.scatter(df_merged["AOD_500nm"], df_merged["AOD_fitted"], alpha=0.7)

mn = min(df_merged["AOD_500nm"].min(), df_merged["AOD_fitted"].min())
mx = max(df_merged["AOD_500nm"].max(), df_merged["AOD_fitted"].max())
plt.plot([mn, mx], [mn, mx], "k--", linewidth=1)

plt.xlabel("Observed AOD$_{500}$")
plt.ylabel("Fitted AOD$_{500}$")
plt.title("Observed vs fitted AOD")
plt.grid(True, alpha=0.3)
plt.tight_layout()

plot1 = os.path.join(out_dir, "Observed_vs_fitted_AOD_multivariate.png")
plt.savefig(plot1, dpi=300, bbox_inches="tight")
plt.show()

print(f"Saved plot: {plot1}")

# ======================================================
# 7. OPTIONAL: TIME SERIES PLOT
# ======================================================

plt.figure(figsize=(10, 5))
plt.plot(df_merged["date"], df_merged["AOD_500nm"], label="Observed AOD", linewidth=1.8)
plt.plot(df_merged["date"], df_merged["AOD_fitted"], label="Fitted AOD", linewidth=1.8)

plt.xlabel("Date")
plt.ylabel("AOD$_{500}$")
plt.title("Observed and fitted AOD time series")
plt.legend()
plt.grid(True, alpha=0.3)
plt.tight_layout()

plot2 = os.path.join(out_dir, "AOD_observed_fitted_timeseries.png")
plt.savefig(plot2, dpi=300, bbox_inches="tight")
plt.show()

print(f"Saved plot: {plot2}")

# =========================================================
# 10. DIRECT RADIATIVE EFFECT / FORCING
# =========================================================

# Constants for DRF calculation
Srad = 461.0
phi = 1.33
Tatm = 0.76
Rs = 0.15
omega = 0.972
beta = 0.21

# Clear-sky and all-sky cloud fraction choices
Cc_clear = 0.0
Cc_allsky = 0.491  # CF calculated from MODIS L2 data

def dre_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)
    )

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

# Extract slope for temperature (ols_slope)
ols_slope = model.params["temp_degC"]

# Calculate DRF values
RE_clear = radiative_efficiency(Cc_clear)
RE_allsky = radiative_efficiency(Cc_allsky)

DRF_dT_clear_ols = RE_clear * ols_slope
DRF_dT_allsky_ols = RE_allsky * ols_slope

# Confidence intervals for the slope
conf_int = model.conf_int().loc["temp_degC"]
conf_int_clear = RE_clear * conf_int
conf_int_allsky = RE_allsky * conf_int

# Save results to a .txt file
output_file = os.path.join(out_dir, "DRF_results.txt")
with open(output_file, "w") as f:
    f.write("Direct Radiative Forcing Results\n")
    f.write("================================\n")
    f.write(f"OLS Slope for Temperature: {ols_slope:.4f}\n")
    f.write(f"Clear-sky Radiative Efficiency: {RE_clear:.4f}\n")
    f.write(f"All-sky Radiative Efficiency: {RE_allsky:.4f}\n")
    f.write(f"DRF (Clear-sky): {DRF_dT_clear_ols:.4f} W/m^2\n")
    f.write(f"DRF (All-sky): {DRF_dT_allsky_ols:.4f} W/m^2\n")
    f.write("\nConfidence Intervals:\n")
    f.write(f"Clear-sky: [{conf_int_clear[0]:.4f}, {conf_int_clear[1]:.4f}] W/m^2\n")
    f.write(f"All-sky: [{conf_int_allsky[0]:.4f}, {conf_int_allsky[1]:.4f}] W/m^2\n")

print(f"Saved DRF results to {output_file}")

# =========================================================
# 11. UPDATED SCATTER PLOT: AOD VS TEMPERATURE
# =========================================================

plt.figure(figsize=(8, 6))

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

# Summer medians as squares
summer_medians = df_merged.groupby(df_merged["date"].dt.month).median(numeric_only=True)
plt.plot(
    summer_medians["temp_degC"].values,
    summer_medians["AOD_500nm"].values,
    marker="s",
    mew=1,
    mec="k",
    markerfacecolor="turquoise",
    markersize=11,
    ls="none",
    label="Summer medians"
)

# Extract x and y values for daily means
x = df_merged["temp_degC"].to_numpy()
y = df_merged["AOD_500nm"].to_numpy()

# Generate x values for the fit line
x_fit = np.linspace(np.min(x), np.max(x), 250)

# Ensure x is a 2D array before adding the constant
if x.ndim == 1:
    x = x.reshape(-1, 1)

# Add constant for intercept
X_daily = sm.add_constant(x, has_constant='add')  # Ensure constant is added

# Define ols_fit_daily by fitting the model
ols_fit_daily = sm.OLS(y, X_daily).fit()

# Extract OLS slope and intercept
ols_intercept = ols_fit_daily.params["const"]
ols_slope = ols_fit_daily.params["x1"]  # x1 corresponds to temp_degC

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

# Compute OLS fit line for daily means
y_fit_ols = ols_slope * x_fit + ols_intercept
pred = ols_fit_daily.get_prediction(sm.add_constant(x_fit))
pred_ci = pred.summary_frame(alpha=0.05)

# Update confidence interval shading
plt.fill_between(
    x_fit,
    pred_ci["mean_ci_lower"].values,
    pred_ci["mean_ci_upper"].values,
    color="lightblue",
    alpha=0.15
)

# Add labels, legend, and grid
plt.xlabel("Temperature (°C)")
plt.ylabel("AOD$_{500}$")
plt.title("Scatter Plot of AOD vs Temperature with Fit and Summer Medians")
plt.legend()
plt.grid(True, alpha=0.3)
plt.tight_layout()

# Save and show the plot
scatter_plot_path = os.path.join(out_dir, "AOD_vs_Temperature_scatterplot_updated.png")
plt.savefig(scatter_plot_path, dpi=300, bbox_inches="tight")
plt.show()

print(f"Saved updated scatter plot: {scatter_plot_path}")

# Debugging: Print model parameters to verify intercept
print("\nModel Parameters:")
print(model.params)

# Ensure constant term is added correctly
if "const" not in X.columns:
    raise ValueError("Constant term (intercept) is missing from the design matrix X.")

# Debugging: Check dimensions of y and X_daily
print(f"Shape of y: {y.shape}")
print(f"Shape of X_daily: {X_daily.shape}")

# Ensure dimensions are compatible
if X_daily.shape[0] != y.shape[0]:
    raise ValueError("Mismatch in dimensions: X_daily and y must have the same number of rows.")

# Debugging: Check if constant is added to X_daily
print("Columns in X_daily:", X_daily.columns)
if "const" not in X_daily.columns:
    raise ValueError("Constant term (intercept) is missing from the design matrix X_daily.")