import numpy as np
import matplotlib.pyplot as plt
import pandas as pd
import os 
from datetime import datetime, timedelta
import fnmatch 
import matplotlib.dates as mdates
import matplotlib.pyplot as plt
from matplotlib.colors import LogNorm
from pathlib import Path
import pytz 
import math
import time
from matplotlib.colors import LinearSegmentedColormap
from matplotlib.colors import ListedColormap, BoundaryNorm
import matplotlib.pyplot as plt
import xarray
import xarray as xr 
import numpy as np
import glob
import getpass
import act 

# Scikit-learn imports
from sklearn.cluster import KMeans
from sklearn.preprocessing import StandardScaler

# Dask imports
from distributed import Client, LocalCluster, progress

# Read in the CSV
df = pd.read_csv("~/DOE_ARM_SummerSchool/merged_cpc_2023_UF_minus_FINE5min.csv")

# Make sure time column is datetime
df["time_Local"] = pd.to_datetime(df["time_Local"])

# Make a date column for grouping by individual days
df["date"] = df["time_Local"].dt.date

# Output file name
output_file = "daily_UF_minus_FINE_stats.csv"
daily_folder = "~/DOE_ARM_SummerSchool/daily_UF_minus_FINE_csvs"

# Define base directory
base_dir = Path("~/DOE_ARM_SummerSchool").expanduser()
# Input CSV
input_file = base_dir / "merged_cpc_2023_UF_minus_FINE5min.csv"
# Read in the CSV
df = pd.read_csv(input_file)
# Make sure time column is datetime
df["time_Local"] = pd.to_datetime(df["time_Local"])
# Make a date column for grouping by individual days
df["date"] = df["time_Local"].dt.date
# Output file name
output_file = "daily_UF_minus_FINE_stats.csv"
# Output folder
daily_folder = base_dir / "daily_UF_minus_FINE_csvs"

# Create/overwrite the output file and write the header
# Create/overwrite output file and write header
with open(output_file, "w") as f:
    f.write(
        "date,"
        "average_UF_minus_FINE,"
        "max_UF_minus_FINE,"
        "min_UF_minus_FINE,"
        "p25_UF_minus_FINE,"
        "p50_UF_minus_FINE,"
        "p75_UF_minus_FINE\n"
    )

# Loop through each individual day
for date, day_df in df.groupby("date"):

    # Copy the day's data so we can safely modify it
    day_df = day_df.copy()

    # For this day, change negative values to NaN
    day_df.loc[day_df["UF_minus_FINE"] < 0, "UF_minus_FINE"] = np.nan

    # Drop NaNs before calculating statistics
    values = day_df["UF_minus_FINE"].dropna()

    # Skip day if no valid data remains
    if len(values) == 0:
        continue

    # Calculate daily statistics
    avg_val = values.mean()
    max_val = values.max()
    min_val = values.min()
    p25 = values.quantile(0.25)
    p50 = values.quantile(0.50)
    p75 = values.quantile(0.75)

    # Append this day's results to the file
    with open(output_file, "a") as f:
        f.write(
            f"{date},"
            f"{avg_val},"
            f"{max_val},"
            f"{min_val},"
            f"{p25},"
            f"{p50},"
            f"{p75}\n"
        )

plot_folder = base_dir / "daily_UF_minus_FINE_plots"
# Loop through each individual day
for date, day_df in df.groupby("date"):
    print(date)
    # Copy the day's data so we can safely modify it
    day_df = day_df.copy()

    # For this day, change negative values to NaN
    day_df.loc[day_df["UF_minus_FINE"] < 0, "UF_minus_FINE"] = np.nan

    # Make sure time column is datetime
    day_df["time_Local"] = pd.to_datetime(day_df["time_Local"])
    
    # Get date from filename or dataframe
    date_str = day_df["time_Local"].dt.date.iloc[0]
    
    plt.figure(figsize=(12, 5))
    plt.plot(day_df["time_Local"], day_df["UF_minus_FINE"])

    plt.title(f"UF minus FINE on {date_str}", fontsize=20)
    plt.xlabel("Time Local", fontsize=16)
    plt.ylabel("UF_minus_FINE", fontsize=16)

    plt.xticks(fontsize=14)
    plt.yticks(fontsize=14)

    plt.grid(True)
    plt.tight_layout()

    # Save plot
    plot_file = os.path.join(plot_folder, f"UF_minus_FINE_{date_str}.png")
    plt.savefig(plot_file, dpi=300)

    # Do not show plot; just close it
    plt.close()