import argparse
import json
import os
import sys
import subprocess
from typing import Dict, Tuple, Any, Optional

import numpy as np
import pandas as pd
import joblib

# ==========================================
# 🛠️ RUNTIME DEPENDENCY INSTALLATION
# ==========================================
try:
    import geopandas as gpd
except ImportError:
    print("📦 Installing geopandas dynamically...", flush=True)
    subprocess.check_call([
        sys.executable, "-m", "pip", "install", "--quiet",
        "geopandas==0.10.2", "shapely<2.0.0", "fiona<1.9.0", "rtree"
    ])
    import geopandas as gpd

DEFAULT_INPUT_DIR = "raw_data"
DEFAULT_OUTPUT_DIR = "processed_data"

KEY_POI_THEMES = {
    "ssot_hawkercentres": "poi_hawker_count",
    "museum": "poi_museum_count",
    "votg_phpc": "poi_clinic_count",
    "libraries": "poi_library_count",
    "communityclubs": "poi_community_club_count",
    "monuments": "poi_monument_count",
    "historicsites": "poi_historic_site_count",
    "hotels": "poi_hotel_count",
    "tourism": "poi_attraction_count",
    "nationalparks": "poi_park_count",
}


def clean_subzone_series(series: pd.Series) -> pd.Series:
    """Normalizes subzone string codes/names for exact joining."""
    return series.astype(str).str.strip().str.upper()


def get_subzone_col(df: pd.DataFrame) -> Optional[str]:
    """Dynamically finds the subzone column name across varying raw CSV source schemas."""
    possible_names = [
        "SUBZONE_N",
        "SUBZONE",
        "SUBZONE_NAME",
        "SZ",
        "SZ_NAME",
        "SUBZONE NAME",
    ]
    for col in df.columns:
        if col.upper().strip() in possible_names or "SUBZONE" in col.upper():
            return col
    return None


def parse_poi_coordinates(
    df_theme: pd.DataFrame,
) -> Optional[gpd.GeoDataFrame]:
    """Parses coordinates from LatLng, Lat/Lng columns, or SVY21 metric columns into EPSG:3414."""
    poi_gdf = None
    col_map = {c.upper().strip(): c for c in df_theme.columns}

    # Strategy 1: LatLng string column parsing
    if "LatLng" in df_theme.columns:
        extracted_floats = (
            df_theme["LatLng"].dropna().astype(str).str.findall(r"(-?\d+\.\d+)")
        )
        valid_pairs = extracted_floats[extracted_floats.str.len() >= 2]

        if not valid_pairs.empty:
            v1 = pd.to_numeric(valid_pairs.str[0], errors="coerce")
            v2 = pd.to_numeric(valid_pairs.str[1], errors="coerce")

            mask1 = (v1 >= 1.1) & (v1 <= 1.5) & (v2 >= 103.5) & (v2 <= 104.1)
            mask2 = (v2 >= 1.1) & (v2 <= 1.5) & (v1 >= 103.5) & (v1 <= 104.1)

            lats, lngs, valid_idx = None, None, None
            if mask1.any():
                lats, lngs, valid_idx = v1[mask1], v2[mask1], v1[mask1].index
            elif mask2.any():
                lats, lngs, valid_idx = v2[mask2], v1[mask2], v2[mask2].index

            if lats is not None and not lats.empty:
                df_filtered = df_theme.loc[valid_idx]
                poi_gdf = gpd.GeoDataFrame(
                    df_filtered,
                    geometry=gpd.points_from_xy(x=lngs, y=lats),
                    crs="EPSG:4326",
                ).to_crs(epsg=3414)

    # Strategy 2: Lat / Long split columns
    if poi_gdf is None:
        lat_key = next(
            (col_map[k] for k in ["LATITUDE", "LAT"] if k in col_map), None
        )
        lng_key = next(
            (
                col_map[k]
                for k in ["LONGITUDE", "LNG", "LON", "LONGTITUDE"]
                if k in col_map
            ),
            None,
        )

        if lat_key and lng_key:
            valid_rows = df_theme.dropna(subset=[lat_key, lng_key]).copy()
            lats = pd.to_numeric(valid_rows[lat_key], errors="coerce")
            lngs = pd.to_numeric(valid_rows[lng_key], errors="coerce")

            valid_mask = (
                (lats >= 1.1)
                & (lats <= 1.5)
                & (lngs >= 103.5)
                & (lngs <= 104.1)
            )

            if valid_mask.any():
                df_filtered = valid_rows.loc[valid_mask]
                poi_gdf = gpd.GeoDataFrame(
                    df_filtered,
                    geometry=gpd.points_from_xy(
                        lngs[valid_mask], lats[valid_mask]
                    ),
                    crs="EPSG:4326",
                ).to_crs(epsg=3414)

    # Strategy 3: Native SVY21 X / Y coordinates
    if poi_gdf is None:
        x_key = next(
            (
                col_map[k]
                for k in ["X", "X_ADDR", "LANDXADDRESSPOINT"]
                if k in col_map
            ),
            None,
        )
        y_key = next(
            (
                col_map[k]
                for k in ["Y", "Y_ADDR", "LANDYADDRESSPOINT"]
                if k in col_map
            ),
            None,
        )

        if x_key and y_key:
            valid_xy = df_theme.dropna(subset=[x_key, y_key]).copy()
            x_vals = pd.to_numeric(valid_xy[x_key], errors="coerce")
            y_vals = pd.to_numeric(valid_xy[y_key], errors="coerce")

            mask_svy21 = (x_vals > 10000) & (y_vals > 10000)
            if mask_svy21.any():
                df_svy21 = valid_xy.loc[mask_svy21]
                poi_gdf = gpd.GeoDataFrame(
                    df_svy21,
                    geometry=gpd.points_from_xy(
                        x_vals[mask_svy21], y_vals[mask_svy21]
                    ),
                    crs="EPSG:3414",
                )

    return poi_gdf


def load_data(input_dir: str) -> Dict[str, Any]:
    """Loads raw vector datasets, CSVs, and micro-themes from input directory."""
    print(f"📥 Loading raw datasets from: {input_dir}", flush=True)
    datasets = {}

    # 1. Base Geometry Frame
    subzone_path = os.path.join(input_dir, "subzone_boundary_mp19.geojson")
    if not os.path.exists(subzone_path):
        raise FileNotFoundError(f"❌ Missing mandatory spatial boundary: {subzone_path}")

    datasets["subzones"] = gpd.read_file(subzone_path).to_crs(epsg=3414)
    datasets["subzones"]["SUBZONE_N"] = clean_subzone_series(datasets["subzones"]["SUBZONE_N"])
    datasets["subzones"]["area_sqkm"] = datasets["subzones"].geometry.area / 1e6

    # 2. Demographics Data
    dwell_path = os.path.join(input_dir, "pop_dwelling_census2020.csv")
    datasets["dwelling"] = pd.read_csv(dwell_path) if os.path.exists(dwell_path) else None

    age_path = os.path.join(input_dir, "pop_age_sex_census2020.csv")
    datasets["age"] = pd.read_csv(age_path) if os.path.exists(age_path) else None

    # 3. Spatial Land Use
    land_path = os.path.join(input_dir, "land_use_mp19.geojson")
    datasets["land_use"] = gpd.read_file(land_path).to_crs(epsg=3414) if os.path.exists(land_path) else None

    # 4. Public Transit Infrastructure & OD Telemetry
    bus_path = os.path.join(input_dir, "lta_bus_stops.csv")
    datasets["bus_stops"] = pd.read_csv(bus_path) if os.path.exists(bus_path) else None

    rail_path = os.path.join(input_dir, "rail_stations_mp19.geojson")
    datasets["rail_stations"] = gpd.read_file(rail_path).to_crs(epsg=3414) if os.path.exists(rail_path) else None

    od_bus_path = os.path.join(input_dir, "lta_od_bus.csv")
    datasets["od_bus"] = pd.read_csv(od_bus_path) if os.path.exists(od_bus_path) else None

    od_train_path = os.path.join(input_dir, "lta_od_train.csv")
    datasets["od_train"] = pd.read_csv(od_train_path) if os.path.exists(od_train_path) else None

    # 5. Education Infrastructure
    school_path = os.path.join(input_dir, "general_info_schools.csv")
    datasets["schools"] = pd.read_csv(school_path) if os.path.exists(school_path) else None

    # 6. OneMap Micro-Themes Data
    datasets["poi_themes"] = {}
    for theme_code in KEY_POI_THEMES.keys():
        theme_file = os.path.join(input_dir, f"onemap_{theme_code}.csv")
        if os.path.exists(theme_file):
            datasets["poi_themes"][theme_code] = pd.read_csv(theme_file)

    return datasets


def build_feature_matrix(input_dir: str) -> pd.DataFrame:
    """Constructs subzone feature matrix by consuming raw datasets."""
    datasets = load_data(input_dir)

    # 1. Base Master Frame
    subzones = datasets["subzones"]
    master_df = subzones[["SUBZONE_N", "PLN_AREA_N", "REGION_N", "area_sqkm"]].copy()

    # 2. Demographics - Housing / Wealth
    dwell_df = datasets.get("dwelling")
    if dwell_df is not None:
        sub_col = get_subzone_col(dwell_df) or "Number"
        dwell_df["SUBZONE_N"] = clean_subzone_series(dwell_df[sub_col])
        dwell_df = dwell_df[
            ~dwell_df["SUBZONE_N"].str.contains(
                "- TOTAL|TOTAL", case=False, na=False
            )
        ].copy()

        num_cols = [
            c
            for c in dwell_df.columns
            if c
            not in [
                "Number",
                "SUBZONE_N",
                "PLN_AREA_N",
                "Planning Area",
                "Total",
            ]
        ]
        for c in num_cols:
            dwell_df[c] = pd.to_numeric(
                dwell_df[c].astype(str).str.replace(",", ""), errors="coerce"
            ).fillna(0)

        condo_cols = [
            c
            for c in num_cols
            if any(
                k in c.upper() for k in ["CONDO", "APARTMENT", "LANDED"]
            )
        ]
        hdb_cols = [c for c in num_cols if "HDB" in c.upper()]

        dwell_df["private_pop"] = (
            dwell_df[condo_cols].sum(axis=1) if condo_cols else 0
        )
        dwell_df["hdb_pop"] = dwell_df[hdb_cols].sum(axis=1) if hdb_cols else 0
        total_dwell = dwell_df["private_pop"] + dwell_df["hdb_pop"]

        dwell_df["private_housing_ratio"] = np.where(
            total_dwell > 0, dwell_df["private_pop"] / total_dwell, 0.0
        )
        dwell_df["hdb_housing_ratio"] = np.where(
            total_dwell > 0, dwell_df["hdb_pop"] / total_dwell, 0.0
        )

        dwell_clean = dwell_df[
            ["SUBZONE_N", "private_housing_ratio", "hdb_housing_ratio"]
        ].drop_duplicates("SUBZONE_N")
        master_df = master_df.merge(
            dwell_clean, on="SUBZONE_N", how="left"
        ).fillna({"private_housing_ratio": 0.0, "hdb_housing_ratio": 0.0})

    # 3. Demographics - Life-stage / Age Ratios
    age_df = datasets.get("age")
    if age_df is not None:
        sub_col = get_subzone_col(age_df) or "Number"
        age_df["SUBZONE_N"] = clean_subzone_series(age_df[sub_col])
        age_df = age_df[
            ~age_df["SUBZONE_N"].str.contains(
                "- TOTAL|TOTAL", case=False, na=False
            )
        ].copy()

        num_cols = [
            c
            for c in age_df.columns
            if c
            not in [
                "Number",
                "SUBZONE_N",
                "PLN_AREA_N",
                "Planning Area",
                "Sex",
                "SEX",
                "Total",
            ]
        ]
        for c in num_cols:
            age_df[c] = pd.to_numeric(
                age_df[c].astype(str).str.replace(",", ""), errors="coerce"
            ).fillna(0)

        age_subzone = age_df.groupby("SUBZONE_N")[num_cols].sum().reset_index()

        young_targets = [
            "0_to_4", "5_to_9", "10_to_14", "0 - 4", "5 - 9", "10 - 14", "0_4", "5_9", "10_14"
        ]
        young_cols = [
            c for c in num_cols if any(t in c.lower() for t in young_targets)
        ]
        elderly_targets = ["65", "70", "75", "80", "85", "90"]
        elderly_cols = [
            c for c in num_cols if any(t in c.lower() for t in elderly_targets)
        ]

        total_pop = age_subzone[num_cols].sum(axis=1)
        age_subzone["young_family_ratio"] = np.where(
            total_pop > 0, age_subzone[young_cols].sum(axis=1) / total_pop, 0.0
        )
        age_subzone["elderly_ratio"] = np.where(
            total_pop > 0, age_subzone[elderly_cols].sum(axis=1) / total_pop, 0.0
        )

        age_clean = age_subzone[
            ["SUBZONE_N", "young_family_ratio", "elderly_ratio"]
        ].drop_duplicates("SUBZONE_N")
        master_df = master_df.merge(
            age_clean, on="SUBZONE_N", how="left"
        ).fillna({"young_family_ratio": 0.0, "elderly_ratio": 0.0})

    # 4. Commercial Land Area
    land_use = datasets.get("land_use")
    if land_use is not None:
        comm_land = land_use[
            land_use["LU_DESC"].isin(
                ["COMMERCIAL", "COMMERCIAL & RESIDENTIAL", "WHITE"]
            )
        ].copy()
        comm_land["comm_area_sqm"] = comm_land.geometry.area

        comm_joined = gpd.sjoin(
            comm_land, subzones[["SUBZONE_N", "geometry"]], predicate="intersects"
        )
        comm_summary = (
            comm_joined.groupby("SUBZONE_N")["comm_area_sqm"]
            .sum()
            .rename("comm_land_sqm")
        )

        master_df = master_df.merge(comm_summary, on="SUBZONE_N", how="left").fillna(
            {"comm_land_sqm": 0.0}
        )
    else:
        master_df["comm_land_sqm"] = 0.0

    master_df["comm_land_ratio"] = master_df["comm_land_sqm"] / (
        master_df["area_sqkm"] * 1e6
    )
    master_df["comm_land_density"] = (
        master_df["comm_land_sqm"] / master_df["area_sqkm"]
    )

    # 5. Transit Nodes & OD Telemetry
    bus_joined, rail_joined = None, None

    bus_stops = datasets.get("bus_stops")
    if bus_stops is not None:
        bus_gdf = gpd.GeoDataFrame(
            bus_stops,
            geometry=gpd.points_from_xy(
                bus_stops["Longitude"], bus_stops["Latitude"]
            ),
            crs="EPSG:4326",
        ).to_crs(epsg=3414)

        bus_joined = gpd.sjoin(
            bus_gdf, subzones[["SUBZONE_N", "geometry"]], predicate="within"
        )
        bus_counts = (
            bus_joined.groupby("SUBZONE_N").size().rename("bus_stop_count")
        )
        master_df = master_df.merge(bus_counts, on="SUBZONE_N", how="left").fillna(
            {"bus_stop_count": 0}
        )
    else:
        master_df["bus_stop_count"] = 0

    rail_stations = datasets.get("rail_stations")
    if rail_stations is not None:
        rail_joined = gpd.sjoin(
            rail_stations,
            subzones[["SUBZONE_N", "geometry"]],
            predicate="intersects",
        )
        rail_counts = (
            rail_joined.groupby("SUBZONE_N").size().rename("mrt_station_count")
        )
        master_df = master_df.merge(
            rail_counts, on="SUBZONE_N", how="left"
        ).fillna({"mrt_station_count": 0})
    else:
        master_df["mrt_station_count"] = 0

    master_df["bus_stop_density"] = (
        master_df["bus_stop_count"] / master_df["area_sqkm"]
    )
    master_df["mrt_station_density"] = (
        master_df["mrt_station_count"] / master_df["area_sqkm"]
    )

    for od_key, name in [
        ("od_bus", "bus_trips"),
        ("od_train", "train_trips"),
    ]:
        od_df = datasets.get(od_key)
        target_joined = bus_joined if "bus" in name else rail_joined

        if (
            od_df is not None
            and target_joined is not None
            and not target_joined.empty
        ):
            trip_cols = [
                c for c in od_df.columns if any(k in c.upper() for k in ["TRIP", "TOTAL", "PASSENGER", "VOL"])
            ]
            dest_cols = [
                c for c in od_df.columns if any(k in c.upper() for k in ["DEST", "PT_CODE", "STATION", "STOP"])
            ]

            if trip_cols and dest_cols:
                trip_col, dest_col = trip_cols[0], dest_cols[0]
                code_col = target_joined.columns[0]

                if "bus" in name and "BusStopCode" in target_joined.columns:
                    code_col = "BusStopCode"
                elif "train" in name:
                    rail_code_cols = [
                        c for c in target_joined.columns if any(k in c.upper() for k in ["STN", "CODE", "STATION"])
                    ]
                    if rail_code_cols:
                        code_col = rail_code_cols[0]

                od_df[dest_col] = (
                    od_df[dest_col].astype(str).str.strip().str.replace(r"\.0$", "", regex=True)
                )
                target_sub = target_joined[[code_col, "SUBZONE_N"]].copy()
                target_sub[code_col] = (
                    target_sub[code_col].astype(str).str.strip().str.replace(r"\.0$", "", regex=True)
                )

                od_merged = od_df.merge(
                    target_sub, left_on=dest_col, right_on=code_col, how="inner"
                )
                vol_summary = (
                    od_merged.groupby("SUBZONE_N")[trip_col]
                    .sum()
                    .rename(f"ambient_{name}")
                )
                master_df = master_df.merge(
                    vol_summary, on="SUBZONE_N", how="left"
                ).fillna({f"ambient_{name}": 0.0})

    if "ambient_bus_trips" not in master_df.columns:
        master_df["ambient_bus_trips"] = 0.0
    if "ambient_train_trips" not in master_df.columns:
        master_df["ambient_train_trips"] = 0.0

    master_df["ambient_bus_trips_density"] = (
        master_df["ambient_bus_trips"] / master_df["area_sqkm"]
    )
    master_df["ambient_train_trips_density"] = (
        master_df["ambient_train_trips"] / master_df["area_sqkm"]
    )

    # 6. Education Infrastructure
    school_df = datasets.get("schools")
    if school_df is not None:
        sub_sch_col = get_subzone_col(school_df)
        if sub_sch_col:
            school_df["SUBZONE_N"] = clean_subzone_series(school_df[sub_sch_col])
            sch_counts = (
                school_df.groupby("SUBZONE_N").size().rename("school_count")
            )
            master_df = master_df.merge(
                sch_counts, on="SUBZONE_N", how="left"
            ).fillna({"school_count": 0})
        else:
            master_df["school_count"] = 0
    else:
        master_df["school_count"] = 0

    master_df["school_density"] = master_df["school_count"] / master_df["area_sqkm"]

    # 7. OneMap Spatial POI Micro-Themes
    subzones_3414 = subzones[["SUBZONE_N", "geometry"]].copy()
    subzones_3414.sindex
    poi_themes = datasets.get("poi_themes", {})

    for theme_code, feature_col in KEY_POI_THEMES.items():
        df_theme = poi_themes.get(theme_code)
        if df_theme is not None and not df_theme.empty:
            poi_gdf = parse_poi_coordinates(df_theme)
            if poi_gdf is not None and not poi_gdf.empty:
                try:
                    joined = gpd.sjoin(
                        poi_gdf,
                        subzones_3414,
                        predicate="intersects",
                        how="inner",
                    )
                    if not joined.empty:
                        joined["SUBZONE_N"] = clean_subzone_series(
                            joined["SUBZONE_N"]
                        )
                        theme_counts = (
                            joined.groupby("SUBZONE_N")
                            .size()
                            .rename(feature_col)
                        )
                        master_df = master_df.merge(
                            theme_counts, on="SUBZONE_N", how="left"
                        ).fillna({feature_col: 0})
                    else:
                        master_df[feature_col] = 0
                except Exception as e:
                    print(f"⚠️ Spatial join failed for theme {theme_code}: {e}", flush=True)
                    master_df[feature_col] = 0
            else:
                master_df[feature_col] = 0
        else:
            master_df[feature_col] = 0

    poi_cols = [c for c in master_df.columns if c.startswith("poi_")]
    master_df["total_micro_poi_density"] = master_df[poi_cols].sum(axis=1) / master_df["area_sqkm"]

    for poi_col in poi_cols:
        master_df[f"{poi_col}_density"] = master_df[poi_col] / master_df["area_sqkm"]

    print(f"⚙️ Raw Feature Matrix assembled: {master_df.shape}", flush=True)
    return master_df


def process_raw_features(
    master_df: pd.DataFrame,
    output_dir: str = DEFAULT_OUTPUT_DIR,
) -> pd.DataFrame:
    """Calculates outlier caps and exports unscaled raw features for model training/Gradia."""
    print("\n" + "=" * 80)
    print("STAGE 2: RAW FEATURE CLEANING & OUTLIER CAPPING SETUP")
    print("=" * 80, flush=True)

    metadata_cols = ["SUBZONE_N", "PLN_AREA_N", "REGION_N"]
    feature_cols = [c for c in master_df.columns if c not in metadata_cols]

    X_raw = master_df[feature_cols].copy()
    X_raw = X_raw.replace([np.inf, -np.inf], np.nan).fillna(0.0)

    # 1. Outlier Clipping (Save caps dictionary for inference)
    print("--> Computing outlier caps (1st-99th percentile)...", flush=True)
    caps = {}
    for col in X_raw.columns:
        lower_cap = float(X_raw[col].quantile(0.01))
        upper_cap = float(X_raw[col].quantile(0.99))
        caps[col] = (lower_cap, upper_cap)

    # 2. Save feature caps artifact for inference.py live transformations
    os.makedirs(output_dir, exist_ok=True)
    joblib.dump(caps, os.path.join(output_dir, "feature_caps.joblib"))
    print(f"💾 Saved feature caps artifact (feature_caps.joblib) to {output_dir}", flush=True)

    # 3. Format unscaled raw dataset
    raw_df = pd.concat([master_df[metadata_cols], X_raw], axis=1)
    print(f"--> Exporting raw feature dataset with {len(feature_cols)} features.\n", flush=True)

    return raw_df


def preprocess_workflow(
    input_dir: str = DEFAULT_INPUT_DIR,
    output_dir: str = DEFAULT_OUTPUT_DIR,
):
    """Orchestrates loading raw inputs, spatial processing, and saving raw features."""
    print("🚀 Initiating Preprocessing Step...", flush=True)

    os.makedirs(output_dir, exist_ok=True)

    # 1. Load raw datasets to grab base boundary layer
    datasets = load_data(input_dir)
    subzones_gdf = datasets["subzones"]

    # 2. Export GeoJSON spatial boundaries to processed output directory
    geojson_output_path = os.path.join(output_dir, "subzone_boundary.geojson")
    subzones_gdf.to_crs(epsg=4326).to_file(geojson_output_path, driver="GeoJSON")
    print(f"🗺️ Spatial boundary GeoJSON saved to: {geojson_output_path}", flush=True)

    # 3. Assemble raw features
    raw_master_df = build_feature_matrix(input_dir=input_dir)
    raw_df = process_raw_features(
        master_df=raw_master_df, 
        output_dir=output_dir
    )

    # 4. Save unscaled raw features matrix
    raw_filepath = os.path.join(output_dir, "raw_features.csv")
    raw_df.to_csv(raw_filepath, index=False)
    print(f"💾 Unscaled raw features matrix saved to: {raw_filepath}", flush=True)


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--input-data", type=str, default=DEFAULT_INPUT_DIR, help="Path to input raw data folder")
    parser.add_argument("--output-data", type=str, default=DEFAULT_OUTPUT_DIR, help="Path to output processed data folder")
    args, _ = parser.parse_known_args()

    preprocess_workflow(
        input_dir=args.input_data,
        output_dir=args.output_data,
    )