import os
import pandas as pd
import numpy as np
import uwmxhelper as uh


def generate_rain_catalog(
    station,
    durations,
    n_extreme_events,
    data_dir=".",
    output_dir="./Mike.model/input",
    min_dry_period=120,
):
    """
    Generate rain event catalogs for one station across multiple durations.

    Parameters
    ----------
    station : str
        MeteoSwiss station code, e.g. "SMA".
    durations : list[int]
        Durations in minutes, e.g. [10, 20, 30, 40, 50, 60].
    n_extreme_events : int
        Number of top extreme windows to keep per duration.
    data_dir : str
        Directory containing MeteoSwiss CSV files.
    swmm_output_dir : str
        Directory where SWMM rainfall files are written.
    min_dry_period : int
        Dry period (minutes) separating events.

    Returns
    -------
    rain_catalog : dict
        Contains counts of events and resulting file paths.
    """

    os.makedirs(output_dir, exist_ok=True)

    print(f"Processing station: {station}")

    rain_catalog_full = []   # collector across ALL durations
    summary = {}

    for duration in durations:
        print(f"Processing duration {duration} min")

        # -------------------------------------
        # Load data
        # -------------------------------------
        
        filename = (
            f"./Data.Rain/Data.Rain.MeteoSwiss/{station}/"
            f"climate-extremes-precipitation-hyeto-data_{duration}-minute-sum_{station}_de.csv"
        )


        if not os.path.exists(filename):
            print(f"[WARN] File not found: {filename}. Skipping duration {duration}.")
            continue

        df = pd.read_csv(filename, sep=";", skiprows=9)
        df["Datum"] = pd.to_datetime(df["Datum"], format="%Y%m%d%H%M")

        # -------------------------------------
        # Preprocessing
        # -------------------------------------
        df = df.sort_values("Datum").drop_duplicates(subset="Datum").reset_index(drop=True)

        # Replace exactly 0.1 mm with 0.0 mm
        df["rain_10min"] = np.where(df["10-Min_Nied"] == 0.1, 0.0, df["10-Min_Nied"])

        # -------------------------------------
        # Event detection
        # -------------------------------------
        step = 10
        dry_steps = min_dry_period // step

        is_wet = df["rain_10min"] > 0.0
        is_dry = ~is_wet

        dry_blocks = (
            is_dry.astype(int)
                  .groupby(is_dry.ne(is_dry.shift()).cumsum())
                  .cumsum()
        )

        event_start = is_wet & (dry_blocks.shift(1) >= dry_steps)
        event_start.iloc[0] = is_wet.iloc[0]

        event_id = event_start.cumsum()
        event_id[is_dry] = np.nan
        df["event_id"] = event_id

        # -------------------------------------
        # Rolling window
        # -------------------------------------
        window_steps = duration // step

        rolled = (
            df.groupby("event_id")["rain_10min"]
              .rolling(window=window_steps, min_periods=window_steps)
              .sum()
              .reset_index(level=0, drop=True)
        )

        df[f"rain_{duration}min"] = rolled

        # -------------------------------------
        # Extract max rainfall window per event
        # -------------------------------------
        records = []

        for eid, ev in df.dropna(subset=["event_id"]).groupby("event_id"):
            series = ev[f"rain_{duration}min"]
            if series.isna().all():
                continue

            idx = series.idxmax()
            if pd.isna(idx):
                continue

            start_idx = idx - (window_steps - 1)
            if start_idx < ev.index.min():
                continue

            records.append({
                "event_id": int(eid),
                "window_start": df.loc[start_idx, "Datum"],
                "window_end": df.loc[idx, "Datum"],
                "max_rain": float(series.max()),
                "start_idx": int(start_idx),
                "end_idx": int(idx),
            })

        events = pd.DataFrame(records)

        # -------------------------------------
        # TOP N EVENTS
        # -------------------------------------
        
        if events.empty:
            print(f"No events found for duration {duration}.")
            continue
        
        top_events = (
            events.sort_values("max_rain", ascending=False)
                  .head(n_extreme_events)
                  .sort_values("window_start")
                  .reset_index(drop=True)
        )
        
        print(f"Selected {len(top_events)} extreme events for duration {duration} min")

        # -------------------------------------
        # Add full events to global catalog
        # -------------------------------------
        selected_eids = top_events["event_id"].unique()

        for eid in selected_eids:
            full_ev = df[df["event_id"] == eid][["Datum", "rain_10min"]]

            for _, r in full_ev.iterrows():
                rain_catalog_full.append({
                    "duration": duration,
                    "event_id": int(eid),
                    "timestamp": r["Datum"],
                    "intensity_10min": r["rain_10min"]
                })

    # ======================================================
    # Save unified rain catalog
    # ======================================================
    rain_catalog_all = (
        pd.DataFrame(rain_catalog_full)
          .drop_duplicates()
          .sort_values(["duration", "event_id", "timestamp"])
          .reset_index(drop=True)
    )

    #out_all = os.path.join(output_dir, f"rain_catalog_full_all_durations_{station}.csv")
    #rain_catalog_all.to_csv(out_all, index=False)

    rain_catalog_model = (
    rain_catalog_all[["timestamp", "intensity_10min"]]
      .drop_duplicates(subset=["timestamp"])
      .sort_values("timestamp")
      .rename(columns={"timestamp": "time", "intensity_10min": "value"})
      .reset_index(drop=True)
)

    out_model = os.path.join(output_dir, f"rain_catalog_{station}.csv")
    rain_catalog_model.to_csv(out_model, index=False)

    # SWMM writer
    swmm_path = os.path.join(output_dir, f"rain_catalog_{station}.dat")
    uh.write_swmm_rainfall(rain_catalog_model, station, swmm_path)

    # ======================================================
    # Summary
    # ======================================================
    total_records = len(rain_catalog_all)
    total_unique_events = rain_catalog_all[["duration", "event_id"]].drop_duplicates().shape[0]

   
    print(f"Final catalog for station {station} saved.")
    print(f"Total unique rain events: {total_unique_events}")

    return {
        "station": station,
        "total_records": total_records,
        "total_unique_events": total_unique_events,
        "model_csv": out_model,
        "swmm_file": swmm_path,
    }