#!/usr/bin/env python3
"""Run the paper's EDA, figures, and classification models.

Each CSV row is treated as an offence or victim record, which matches the unit
in the source data. OFFENCE and UCR fields are not model predictors because they
reveal the target category.
"""

from __future__ import annotations

import argparse
import json
import math
import warnings
from pathlib import Path

import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt
from matplotlib.colors import LogNorm
import numpy as np
import pandas as pd
from sklearn.compose import ColumnTransformer
from sklearn.dummy import DummyClassifier
from sklearn.ensemble import RandomForestClassifier
from sklearn.impute import SimpleImputer
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import (
    ConfusionMatrixDisplay,
    accuracy_score,
    balanced_accuracy_score,
    classification_report,
    confusion_matrix,
    f1_score,
)
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import OneHotEncoder, OrdinalEncoder, StandardScaler


RANDOM_SEED = 5430
ANALYSIS_START_YEAR = 2014
ANALYSIS_END_YEAR = 2025
TRAIN_END_YEAR = 2023
VALIDATION_YEAR = 2024
TEST_YEAR = 2025

CATEGORY_ORDER = [
    "Assault",
    "Break and Enter",
    "Auto Theft",
    "Robbery",
    "Theft Over",
]

CATEGORY_COLORS = {
    "Assault": "#355C7D",
    "Break and Enter": "#C06C84",
    "Auto Theft": "#F8B195",
    "Robbery": "#6C5B7B",
    "Theft Over": "#2A9D8F",
}

WEEKDAY_ORDER = [
    "Monday",
    "Tuesday",
    "Wednesday",
    "Thursday",
    "Friday",
    "Saturday",
    "Sunday",
]

REQUIRED_COLUMNS = [
    "EVENT_UNIQUE_ID",
    "REPORT_YEAR",
    "OCC_YEAR",
    "OCC_MONTH",
    "OCC_DOY",
    "OCC_DOW",
    "OCC_HOUR",
    "DIVISION",
    "LOCATION_TYPE",
    "PREMISES_TYPE",
    "CSI_CATEGORY",
    "NEIGHBOURHOOD_158",
    "LONG_WGS84",
    "LAT_WGS84",
]

CATEGORICAL_FEATURES = [
    "OCC_MONTH",
    "OCC_DOW",
    "PREMISES_TYPE",
    "LOCATION_TYPE",
    "DIVISION",
    "NEIGHBOURHOOD_158",
]

NUMERIC_FEATURES = [
    "hour_sin",
    "hour_cos",
    "doy_sin",
    "doy_cos",
    "LONG_WGS84",
    "LAT_WGS84",
]


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--input", type=Path, required=True, help="Path to the CSI CSV")
    parser.add_argument(
        "--output",
        type=Path,
        default=Path("output"),
        help="Directory for generated files",
    )
    parser.add_argument(
        "--sample-output",
        type=Path,
        help="Write a fixed-seed 5,000-row sample to this path",
    )
    parser.add_argument(
        "--max-train-rows",
        type=int,
        default=None,
        help="Limit training rows with fixed-seed sampling",
    )
    parser.add_argument(
        "--skip-model",
        action="store_true",
        help="Stop after creating the EDA figures",
    )
    return parser.parse_args()


def load_data(path: Path) -> pd.DataFrame:
    if not path.exists():
        raise FileNotFoundError(f"Input CSV not found: {path}")
    header = pd.read_csv(path, nrows=0).columns.tolist()
    missing = sorted(set(REQUIRED_COLUMNS) - set(header))
    if missing:
        raise ValueError(f"Input is missing required columns: {missing}")
    df = pd.read_csv(path, usecols=REQUIRED_COLUMNS, low_memory=False)
    for column in CATEGORICAL_FEATURES + ["CSI_CATEGORY"]:
        df[column] = df[column].astype("string").str.strip().fillna("Missing")
    for column in ["REPORT_YEAR", "OCC_YEAR", "OCC_DOY", "OCC_HOUR"]:
        df[column] = pd.to_numeric(df[column], errors="coerce")
    for column in ["LONG_WGS84", "LAT_WGS84"]:
        df[column] = pd.to_numeric(df[column], errors="coerce")
    return df


def analysis_frame(df: pd.DataFrame) -> pd.DataFrame:
    mask = (
        df["REPORT_YEAR"].between(ANALYSIS_START_YEAR, ANALYSIS_END_YEAR)
        & df["CSI_CATEGORY"].isin(CATEGORY_ORDER)
    )
    return df.loc[mask].copy()


def add_cyclic_features(df: pd.DataFrame) -> pd.DataFrame:
    result = df.copy()
    hour = result["OCC_HOUR"].fillna(0).clip(0, 23)
    day = result["OCC_DOY"].fillna(1).clip(1, 366)
    result["hour_sin"] = np.sin(2 * np.pi * hour / 24)
    result["hour_cos"] = np.cos(2 * np.pi * hour / 24)
    result["doy_sin"] = np.sin(2 * np.pi * (day - 1) / 366)
    result["doy_cos"] = np.cos(2 * np.pi * (day - 1) / 366)
    valid_geo = result["LONG_WGS84"].between(-80.1, -78.8) & result[
        "LAT_WGS84"
    ].between(43.4, 44.0)
    result.loc[~valid_geo, ["LONG_WGS84", "LAT_WGS84"]] = np.nan
    return result


def create_sample(df: pd.DataFrame, path: Path, n: int = 5000) -> None:
    """Save a fixed-seed sample stratified by report year and category."""
    frame = analysis_frame(df)
    if len(frame) <= n:
        sampled = frame.copy()
    else:
        fractions = frame.groupby(["REPORT_YEAR", "CSI_CATEGORY"], observed=True)[
            "EVENT_UNIQUE_ID"
        ].transform("size")
        target = np.maximum(2, np.round(n * fractions / len(frame))).astype(int)
        frame = frame.assign(_target=target)
        parts = []
        for _, group in frame.groupby(
            ["REPORT_YEAR", "CSI_CATEGORY"], sort=True, observed=True
        ):
            parts.append(
                group.sample(
                    n=min(int(group["_target"].iloc[0]), len(group)),
                    random_state=RANDOM_SEED,
                )
            )
        sampled = pd.concat(parts).drop(columns="_target")
        if len(sampled) > n:
            sampled = sampled.sample(n=n, random_state=RANDOM_SEED)
    path.parent.mkdir(parents=True, exist_ok=True)
    sampled.sort_values(["REPORT_YEAR", "EVENT_UNIQUE_ID"]).to_csv(path, index=False)


def dataset_summary(raw: pd.DataFrame, frame: pd.DataFrame) -> dict:
    valid_geo = raw["LONG_WGS84"].between(-80.1, -78.8) & raw[
        "LAT_WGS84"
    ].between(43.4, 44.0)
    event_categories = raw.groupby("EVENT_UNIQUE_ID", observed=True)[
        "CSI_CATEGORY"
    ].nunique()
    counts = frame["CSI_CATEGORY"].value_counts().reindex(CATEGORY_ORDER)
    annual = frame.groupby("REPORT_YEAR", observed=True).size()
    return {
        "raw_rows": int(len(raw)),
        "analysis_rows_2014_2025": int(len(frame)),
        "unique_event_ids": int(raw["EVENT_UNIQUE_ID"].nunique()),
        "additional_offence_or_victim_rows": int(
            len(raw) - raw["EVENT_UNIQUE_ID"].nunique()
        ),
        "events_with_multiple_csi_categories": int((event_categories > 1).sum()),
        "valid_coordinate_rows": int(valid_geo.sum()),
        "invalid_or_suppressed_coordinate_rows": int((~valid_geo).sum()),
        "category_counts_2014_2025": {k: int(v) for k, v in counts.items()},
        "category_percent_2014_2025": {
            k: round(100 * int(v) / len(frame), 2) for k, v in counts.items()
        },
        "annual_counts": {str(int(k)): int(v) for k, v in annual.items()},
        "partial_2026_rows_excluded": int((raw["REPORT_YEAR"] == 2026).sum()),
    }


def set_plot_style() -> None:
    plt.rcParams.update(
        {
            "font.family": "DejaVu Sans",
            "font.size": 8.5,
            "axes.titlesize": 9.5,
            "axes.labelsize": 8.5,
            "axes.spines.top": False,
            "axes.spines.right": False,
            "figure.dpi": 150,
            "savefig.dpi": 300,
            "savefig.bbox": "tight",
            "pdf.fonttype": 42,
            "ps.fonttype": 42,
        }
    )


def save_figure(fig: plt.Figure, stem: Path) -> None:
    stem.parent.mkdir(parents=True, exist_ok=True)
    fig.savefig(stem.with_suffix(".pdf"))
    fig.savefig(stem.with_suffix(".png"))
    plt.close(fig)


def plot_eda(frame: pd.DataFrame, figures_dir: Path) -> None:
    set_plot_style()
    fig, axes = plt.subplots(2, 2, figsize=(7.1, 5.3), constrained_layout=True)

    counts = frame["CSI_CATEGORY"].value_counts().reindex(CATEGORY_ORDER)
    bars = axes[0, 0].barh(
        CATEGORY_ORDER,
        counts.values / 1000,
        color=[CATEGORY_COLORS[c] for c in CATEGORY_ORDER],
    )
    axes[0, 0].invert_yaxis()
    axes[0, 0].set_xlabel("Records (thousands)")
    axes[0, 0].set_title("(a) Class distribution")
    for bar, count in zip(bars, counts.values):
        axes[0, 0].text(
            bar.get_width() + 1,
            bar.get_y() + bar.get_height() / 2,
            f"{100 * count / len(frame):.1f}%",
            va="center",
            fontsize=7.5,
        )

    annual = (
        frame.groupby(["REPORT_YEAR", "CSI_CATEGORY"], observed=True)
        .size()
        .unstack(fill_value=0)
        .reindex(columns=CATEGORY_ORDER)
    )
    for category in CATEGORY_ORDER:
        axes[0, 1].plot(
            annual.index,
            annual[category] / 1000,
            marker="o",
            markersize=2.6,
            linewidth=1.25,
            color=CATEGORY_COLORS[category],
            label=category,
        )
    axes[0, 1].set_title("(b) Annual reported records")
    axes[0, 1].set_ylabel("Records (thousands)")
    axes[0, 1].set_xticks([2014, 2016, 2018, 2020, 2022, 2024, 2025])
    axes[0, 1].tick_params(axis="x", rotation=35)
    axes[0, 1].grid(axis="y", alpha=0.25)
    axes[0, 1].legend(fontsize=6.2, ncol=2, frameon=False, loc="upper left")

    hourly = (
        frame.groupby(["OCC_HOUR", "CSI_CATEGORY"], observed=True)
        .size()
        .unstack(fill_value=0)
        .reindex(index=range(24), fill_value=0)
        .reindex(columns=CATEGORY_ORDER)
    )
    hourly = hourly.div(hourly.sum(axis=0), axis=1) * 100
    for category in CATEGORY_ORDER:
        axes[1, 0].plot(
            hourly.index,
            hourly[category],
            linewidth=1.35,
            color=CATEGORY_COLORS[category],
            label=category,
        )
    axes[1, 0].set_title("(c) Within-class hourly profile")
    axes[1, 0].set_xlabel("Occurrence hour")
    axes[1, 0].set_ylabel("Share of class records (%)")
    axes[1, 0].set_xticks(range(0, 24, 3))
    axes[1, 0].grid(alpha=0.2)

    premise = pd.crosstab(
        frame["PREMISES_TYPE"], frame["CSI_CATEGORY"], normalize="columns"
    ).reindex(columns=CATEGORY_ORDER)
    premise = premise.drop(index="Missing", errors="ignore") * 100
    image = axes[1, 1].imshow(premise.values, cmap="Blues", aspect="auto", vmin=0)
    axes[1, 1].set_title("(d) Premises composition within class")
    axes[1, 1].set_xticks(range(len(CATEGORY_ORDER)))
    axes[1, 1].set_xticklabels(
        ["Assault", "B&E", "Auto\nTheft", "Robbery", "Theft\nOver"],
        rotation=0,
    )
    axes[1, 1].set_yticks(range(len(premise.index)))
    axes[1, 1].set_yticklabels(premise.index)
    for i in range(len(premise.index)):
        for j in range(len(CATEGORY_ORDER)):
            value = premise.iloc[i, j]
            axes[1, 1].text(
                j,
                i,
                f"{value:.0f}",
                ha="center",
                va="center",
                fontsize=6.8,
                color="white" if value > premise.values.max() * 0.55 else "#222222",
            )
    colorbar = fig.colorbar(image, ax=axes[1, 1], fraction=0.046, pad=0.03)
    colorbar.set_label("Percent")

    save_figure(fig, figures_dir / "fig1_eda_overview")


def plot_weekday_profile(frame: pd.DataFrame, figures_dir: Path) -> None:
    set_plot_style()
    weekday = pd.crosstab(
        frame["CSI_CATEGORY"], frame["OCC_DOW"], normalize="index"
    ).reindex(index=CATEGORY_ORDER, columns=WEEKDAY_ORDER, fill_value=0)
    weekday = weekday * 100

    fig, ax = plt.subplots(figsize=(3.45, 2.05), constrained_layout=True)
    image = ax.imshow(
        weekday.values,
        cmap="Blues",
        aspect="auto",
        vmin=0,
        vmax=18,
    )
    ax.set_title("Within-class weekday profile")
    ax.set_xticks(range(len(WEEKDAY_ORDER)))
    ax.set_xticklabels(["Mon", "Tue", "Wed", "Thu", "Fri", "Sat", "Sun"])
    ax.set_yticks(range(len(CATEGORY_ORDER)))
    ax.set_yticklabels(["Assault", "B&E", "Auto Theft", "Robbery", "Theft Over"])
    ax.tick_params(length=0, labelsize=7)

    for i in range(len(CATEGORY_ORDER)):
        for j in range(len(WEEKDAY_ORDER)):
            value = weekday.iloc[i, j]
            ax.text(
                j,
                i,
                f"{value:.1f}",
                ha="center",
                va="center",
                fontsize=6.2,
                color="white" if value >= 9 else "#222222",
            )

    colorbar = fig.colorbar(image, ax=ax, fraction=0.055, pad=0.025)
    colorbar.set_label("Percent")
    save_figure(fig, figures_dir / "fig_weekday_profile")


def plot_spatial(frame: pd.DataFrame, figures_dir: Path) -> None:
    geo = frame[
        frame["LONG_WGS84"].between(-79.75, -79.10)
        & frame["LAT_WGS84"].between(43.55, 43.90)
    ].copy()
    panels = ["All CSI records"] + CATEGORY_ORDER
    fig, axes = plt.subplots(2, 3, figsize=(7.1, 4.7), constrained_layout=True)
    for ax, panel in zip(axes.flat, panels):
        data = geo if panel == "All CSI records" else geo[geo["CSI_CATEGORY"] == panel]
        ax.hexbin(
            data["LONG_WGS84"],
            data["LAT_WGS84"],
            gridsize=52,
            mincnt=1,
            cmap="magma_r",
            norm=LogNorm(),
            linewidths=0,
        )
        ax.set_xlim(-79.68, -79.15)
        ax.set_ylim(43.57, 43.86)
        ax.set_aspect(1.35)
        ax.set_title(panel)
        ax.tick_params(labelsize=6.5)
        ax.set_xlabel("Longitude")
        ax.set_ylabel("Latitude")
        ax.grid(alpha=0.12)
    fig.suptitle(
        "Spatial density of geocoded offence/victim records, 2014 to 2025",
        fontsize=10,
    )
    save_figure(fig, figures_dir / "fig2_spatial_density")


def prepare_model_data(frame: pd.DataFrame) -> tuple[pd.DataFrame, pd.Series]:
    model = add_cyclic_features(frame)
    return model[CATEGORICAL_FEATURES + NUMERIC_FEATURES], model["CSI_CATEGORY"]


def metric_row(name: str, y_true: pd.Series, prediction: np.ndarray) -> dict:
    return {
        "model": name,
        "accuracy": float(accuracy_score(y_true, prediction)),
        "balanced_accuracy": float(balanced_accuracy_score(y_true, prediction)),
        "macro_f1": float(f1_score(y_true, prediction, average="macro")),
        "weighted_f1": float(f1_score(y_true, prediction, average="weighted")),
    }


def model_analysis(
    frame: pd.DataFrame, output_dir: Path, max_train_rows: int | None
) -> list[dict]:
    train = frame[frame["REPORT_YEAR"] <= TRAIN_END_YEAR].copy()
    validation = frame[frame["REPORT_YEAR"] == VALIDATION_YEAR].copy()
    test = frame[frame["REPORT_YEAR"] == TEST_YEAR].copy()
    if min(len(train), len(validation), len(test)) == 0:
        raise ValueError(
            "Modeling requires records in training years through 2023, 2024, and 2025."
        )
    if max_train_rows and len(train) > max_train_rows:
        train = train.sample(n=max_train_rows, random_state=RANDOM_SEED)

    x_train, y_train = prepare_model_data(train)
    x_validation, y_validation = prepare_model_data(validation)
    x_test, y_test = prepare_model_data(test)

    numeric_pipe = Pipeline(
        [
            ("imputer", SimpleImputer(strategy="median")),
            ("scale", StandardScaler()),
        ]
    )
    categorical_pipe = Pipeline(
        [
            ("imputer", SimpleImputer(strategy="most_frequent")),
            (
                "onehot",
                OneHotEncoder(handle_unknown="ignore", min_frequency=10),
            ),
        ]
    )
    logistic_preprocessor = ColumnTransformer(
        [
            ("numeric", numeric_pipe, NUMERIC_FEATURES),
            ("categorical", categorical_pipe, CATEGORICAL_FEATURES),
        ]
    )
    logistic = Pipeline(
        [
            ("preprocess", logistic_preprocessor),
            (
                "model",
                LogisticRegression(
                    max_iter=250,
                    tol=1e-3,
                    solver="lbfgs",
                    class_weight="balanced",
                    random_state=RANDOM_SEED,
                ),
            ),
        ]
    )

    ordinal_preprocessor = ColumnTransformer(
        [
            (
                "categorical",
                OrdinalEncoder(
                    handle_unknown="use_encoded_value",
                    unknown_value=-1,
                    encoded_missing_value=-1,
                ),
                CATEGORICAL_FEATURES,
            ),
            ("numeric", SimpleImputer(strategy="median"), NUMERIC_FEATURES),
        ]
    )
    forest = Pipeline(
        [
            ("preprocess", ordinal_preprocessor),
            (
                "model",
                RandomForestClassifier(
                    n_estimators=220,
                    max_depth=24,
                    min_samples_leaf=4,
                    max_features="sqrt",
                    class_weight="balanced_subsample",
                    n_jobs=-1,
                    random_state=RANDOM_SEED,
                ),
            ),
        ]
    )

    dummy = DummyClassifier(strategy="most_frequent")
    dummy.fit(np.zeros((len(y_train), 1)), y_train)
    metrics = []
    predictions: dict[str, np.ndarray] = {}
    metrics.append(
        metric_row(
            "Majority baseline",
            y_test,
            dummy.predict(np.zeros((len(y_test), 1))),
        )
    )

    validation_scores = {}
    fitted_models = {}
    for name, model in [("Logistic regression", logistic), ("Random forest", forest)]:
        model.fit(x_train, y_train)
        fitted_models[name] = model
        validation_prediction = model.predict(x_validation)
        validation_scores[name] = metric_row(
            name, y_validation, validation_prediction
        )
        prediction = model.predict(x_test)
        predictions[name] = prediction
        metrics.append(metric_row(name, y_test, prediction))
        report = classification_report(
            y_test,
            prediction,
            labels=CATEGORY_ORDER,
            output_dict=True,
            zero_division=0,
        )
        report_path = output_dir / f"classification_report_{name.lower().replace(' ', '_')}.json"
        report_path.write_text(json.dumps(report, indent=2), encoding="utf-8")

    best_name = max(validation_scores, key=lambda name: validation_scores[name]["macro_f1"])
    best_prediction = predictions[best_name]
    best_model = fitted_models[best_name]

    details = {
        "split": {
            "train_years": f"{ANALYSIS_START_YEAR}-{TRAIN_END_YEAR}",
            "validation_year": VALIDATION_YEAR,
            "test_year": TEST_YEAR,
            "train_rows": int(len(train)),
            "validation_rows": int(len(validation)),
            "test_rows": int(len(test)),
        },
        "selection_metric": "validation macro-F1",
        "validation_metrics": validation_scores,
        "selected_model": best_name,
        "test_metrics": metrics,
    }
    (output_dir / "model_metrics.json").write_text(
        json.dumps(details, indent=2), encoding="utf-8"
    )
    pd.DataFrame(metrics).to_csv(output_dir / "model_metrics.csv", index=False)

    plot_model_results(
        y_test,
        best_prediction,
        best_name,
        best_model,
        output_dir / "figures",
    )
    return metrics


def plot_model_results(
    y_true: pd.Series,
    prediction: np.ndarray,
    model_name: str,
    fitted_model: Pipeline,
    figures_dir: Path,
) -> None:
    set_plot_style()
    fig, axes = plt.subplots(1, 2, figsize=(7.1, 2.7), constrained_layout=True)
    matrix = confusion_matrix(
        y_true, prediction, labels=CATEGORY_ORDER, normalize="true"
    )
    display = ConfusionMatrixDisplay(
        confusion_matrix=matrix,
        display_labels=["Assault", "B&E", "Auto Theft", "Robbery", "Theft Over"],
    )
    display.plot(ax=axes[0], cmap="Blues", colorbar=False, values_format=".2f")
    axes[0].set_title(f"(a) 2025 normalized confusion matrix\n{model_name}")
    axes[0].tick_params(axis="x", rotation=35, labelsize=6.5)
    axes[0].tick_params(axis="y", labelsize=6.5)

    if model_name == "Random forest":
        importances = fitted_model.named_steps["model"].feature_importances_
        feature_names = CATEGORICAL_FEATURES + NUMERIC_FEATURES
        order = np.argsort(importances)
        axes[1].barh(
            np.array(feature_names)[order], importances[order], color="#355C7D"
        )
        axes[1].set_xlabel("Mean decrease in impurity")
        axes[1].set_title("(b) Random-forest feature importance")
    else:
        report = classification_report(
            y_true,
            prediction,
            labels=CATEGORY_ORDER,
            output_dict=True,
            zero_division=0,
        )
        f1_values = [report[c]["f1-score"] for c in CATEGORY_ORDER]
        axes[1].barh(
            CATEGORY_ORDER,
            f1_values,
            color=[CATEGORY_COLORS[c] for c in CATEGORY_ORDER],
        )
        axes[1].invert_yaxis()
        axes[1].set_xlim(0, 1)
        axes[1].set_xlabel("F1-score")
        axes[1].set_title("(b) 2025 class-specific F1-score")
    save_figure(fig, figures_dir / "fig3_model_results")


def main() -> None:
    args = parse_args()
    args.output.mkdir(parents=True, exist_ok=True)
    (args.output / "figures").mkdir(parents=True, exist_ok=True)
    raw = load_data(args.input)
    frame = analysis_frame(raw)
    summary = dataset_summary(raw, frame)
    (args.output / "dataset_summary.json").write_text(
        json.dumps(summary, indent=2), encoding="utf-8"
    )
    if args.sample_output:
        create_sample(raw, args.sample_output)
    plot_eda(frame, args.output / "figures")
    plot_weekday_profile(frame, args.output / "figures")
    plot_spatial(frame, args.output / "figures")
    if not args.skip_model:
        with warnings.catch_warnings():
            warnings.simplefilter("ignore", category=FutureWarning)
            model_analysis(frame, args.output, args.max_train_rows)
    print(f"Finished. Files are in {args.output}")


if __name__ == "__main__":
    main()
