from __future__ import annotations

from collections import Counter, defaultdict
from datetime import datetime

import numpy as np
import pandas as pd
from lightfm import LightFM
from lightfm.cross_validation import random_train_test_split
from lightfm.data import Dataset
from lightfm.evaluation import auc_score, precision_at_k, recall_at_k
from tqdm.auto import tqdm

from .console import console
from .config import ProjectPaths, RuntimeSettings, TrainingConfig
from .io import dump_pickle, read_csv, write_csv, write_json


def normalize_status(status) -> str:
    if not isinstance(status, str):
        return ""
    return status.strip().lower().replace("-", "_").replace(" ", "_")


def interaction_weight(status: str, score: float) -> float:
    normalized_status = normalize_status(status)
    numeric_score = float(score or 0)

    if normalized_status in {"completed", "watching"}:
        if numeric_score == 0:
            return 0.5
        if numeric_score >= 8:
            return numeric_score / 10.0
        if numeric_score == 7:
            return 0.6
        if numeric_score == 6:
            return 0.4
        if numeric_score == 5:
            return 0.2
        return 0.1
    if normalized_status == "plan_to_watch":
        return 0.7
    if normalized_status == "dropped":
        return 0.0
    if normalized_status in {"on_hold", "hold"}:
        return 0.2
    return 0.1


def filter_interactions(
    interactions_df: pd.DataFrame,
    *,
    min_interactions_per_user: int,
    min_interactions_per_item: int,
) -> pd.DataFrame:
    filtered_df = interactions_df.copy()

    while True:
        user_counts = filtered_df["username"].value_counts()
        item_counts = filtered_df["anime_id"].value_counts()
        initial_rows = len(filtered_df)

        filtered_df = filtered_df[
            filtered_df["username"].isin(
                user_counts[user_counts >= min_interactions_per_user].index
            )
        ]
        filtered_df = filtered_df[
            filtered_df["anime_id"].isin(
                item_counts[item_counts >= min_interactions_per_item].index
            )
        ]

        if len(filtered_df) == initial_rows:
            break

    return filtered_df


def derive_user_genre_features(
    interactions_df: pd.DataFrame,
    catalog_df: pd.DataFrame,
    *,
    output_path=None,
) -> pd.DataFrame:
    demographics = set(RuntimeSettings().demographic_genres)

    catalog_lookup = catalog_df[["anime_id", "genres"]].copy()
    catalog_lookup["anime_id"] = catalog_lookup["anime_id"].astype(str)

    merged_df = interactions_df.merge(catalog_lookup, on="anime_id", how="left")
    total_counts: Counter[str] = Counter()
    genre_counts: dict[str, Counter[str]] = defaultdict(Counter)

    with tqdm(total=len(merged_df), desc="Deriving user genres", unit="row", dynamic_ncols=True) as bar:
        for idx, row in enumerate(merged_df.itertuples(index=False), start=1):
            username = str(row.username)
            genres_raw = getattr(row, "genres", "")
            genres_str = str(genres_raw) if pd.notna(genres_raw) else ""
            total_counts[username] += 1

            parsed_genres = {
                genre.strip()
                for genre in genres_str.split(",")
                if genre.strip() and genre.strip() not in demographics
            }
            for genre in parsed_genres:
                genre_counts[username][genre] += 1

            if idx % 5000 == 0 or idx == len(merged_df):
                bar.update(idx - bar.n)
                bar.set_postfix(users=len(total_counts), genres=sum(len(c) for c in genre_counts.values()))

    rows: list[dict] = []
    for username, counts in genre_counts.items():
        denominator = total_counts.get(username, 0)
        if denominator == 0:
            continue
        for genre, count in counts.items():
            rows.append(
                {
                    "username": username,
                    "genre": genre,
                    "weight": round(count / denominator, 6),
                }
            )

    feature_df = pd.DataFrame(rows, columns=["username", "genre", "weight"])
    if not feature_df.empty:
        feature_df = feature_df.sort_values(["username", "genre"])
    if output_path is not None:
        write_csv(feature_df, output_path)
    return feature_df


def load_user_genre_features(feature_df: pd.DataFrame) -> list[tuple[str, dict[str, float]]]:
    if feature_df.empty:
        return []

    feature_map: dict[str, dict[str, float]] = defaultdict(dict)
    for row in feature_df.itertuples(index=False):
        feature_map[str(row.username)][str(row.genre)] = float(row.weight)
    return sorted(feature_map.items(), key=lambda item: item[0])


def train_model(
    paths: ProjectPaths | None = None,
    config: TrainingConfig | None = None,
) -> dict:
    paths = paths or ProjectPaths()
    config = config or TrainingConfig()
    paths.ensure_directories()
    console.info("Loading interaction data and anime catalog for training")

    interactions_path = paths.resolve_user_interactions_input()
    catalog_path = paths.resolve_anime_catalog_input()

    interactions_df = read_csv(interactions_path)
    interactions_df["anime_id"] = interactions_df["anime_id"].astype(str)
    interactions_df["username"] = interactions_df["username"].astype(str)
    console.info(
        f"Loaded {len(interactions_df)} interactions across {interactions_df['username'].nunique()} users"
    )

    filtered_df = filter_interactions(
        interactions_df,
        min_interactions_per_user=config.min_interactions_per_user,
        min_interactions_per_item=config.min_interactions_per_item,
    )
    if filtered_df.empty:
        raise RuntimeError("No interactions remain after filtering. Lower the thresholds.")
    console.info(
        "After filtering: "
        f"{len(filtered_df)} interactions, "
        f"{filtered_df['username'].nunique()} users, "
        f"{filtered_df['anime_id'].nunique()} anime"
    )

    catalog_df = read_csv(catalog_path)
    catalog_df["anime_id"] = catalog_df["anime_id"].astype(str)
    console.info(f"Loaded anime catalog with {len(catalog_df)} rows")

    feature_df = derive_user_genre_features(
        filtered_df[["username", "anime_id"]].copy(),
        catalog_df[["anime_id", "genres"]].copy(),
        output_path=paths.user_genre_features,
    )
    console.info(
        f"Derived {len(feature_df)} user-genre feature rows into {paths.user_genre_features}"
    )

    user_features_input = load_user_genre_features(feature_df)
    active_users = set(filtered_df["username"].unique())
    filtered_user_features_input = [
        (username, features)
        for username, features in user_features_input
        if username in active_users
    ]

    all_genres = sorted(
        {
            genre
            for _, features in filtered_user_features_input
            for genre in features.keys()
        }
    )

    dataset = Dataset()
    dataset.fit(
        users=filtered_df["username"].unique(),
        items=filtered_df["anime_id"].unique(),
        user_features=all_genres,
    )
    console.info(
        f"LightFM dataset fit with {filtered_df['username'].nunique()} users, "
        f"{filtered_df['anime_id'].nunique()} anime, {len(all_genres)} genre features"
    )

    interactions_data = [
        (row.username, row.anime_id, interaction_weight(row.status, row.score))
        for row in filtered_df.itertuples(index=False)
    ]
    interactions, weights = dataset.build_interactions(interactions_data)
    user_features_matrix = dataset.build_user_features(
        filtered_user_features_input,
        normalize=False,
    )
    console.info(
        f"Matrix shapes | interactions={interactions.shape} | user_features={user_features_matrix.shape}"
    )

    random_state = np.random.RandomState(config.random_seed)
    train_interactions, test_interactions = random_train_test_split(
        interactions,
        test_percentage=config.test_set_fraction,
        random_state=random_state,
    )
    train_weights, _ = random_train_test_split(
        weights,
        test_percentage=config.test_set_fraction,
        random_state=np.random.RandomState(config.random_seed),
    )

    model = LightFM(
        no_components=config.n_components,
        learning_rate=config.learning_rate,
        loss=config.loss_function,
        random_state=config.random_seed,
    )
    console.info(
        "Training LightFM model "
        f"(components={config.n_components}, epochs={config.epochs}, loss={config.loss_function})"
    )
    model.fit(
        train_interactions,
        user_features=user_features_matrix,
        sample_weight=train_weights,
        epochs=config.epochs,
        num_threads=config.threads,
        verbose=True,
    )

    dump_pickle(model, paths.model, paths.legacy_model)
    dump_pickle(dataset, paths.dataset, paths.legacy_dataset)
    console.success(f"Saved model to {paths.model}")
    console.success(f"Saved dataset to {paths.dataset}")

    metrics = {}
    if test_interactions is not None and test_interactions.nnz > 0:
        metrics = {
            "precision_at_10": float(
                precision_at_k(
                    model,
                    test_interactions,
                    train_interactions=train_interactions,
                    user_features=user_features_matrix,
                    k=10,
                    num_threads=config.threads,
                ).mean()
            ),
            "recall_at_10": float(
                recall_at_k(
                    model,
                    test_interactions,
                    train_interactions=train_interactions,
                    user_features=user_features_matrix,
                    k=10,
                    num_threads=config.threads,
                ).mean()
            ),
            "auc": float(
                auc_score(
                    model,
                    test_interactions,
                    train_interactions=train_interactions,
                    user_features=user_features_matrix,
                    num_threads=config.threads,
                ).mean()
            ),
        }

    summary = {
        "generated_at": datetime.utcnow().isoformat() + "Z",
        "interactions_path": str(interactions_path),
        "catalog_path": str(catalog_path),
        "feature_path": str(paths.user_genre_features),
        "model_path": str(paths.model),
        "dataset_path": str(paths.dataset),
        "rows_before_filtering": int(len(interactions_df)),
        "rows_after_filtering": int(len(filtered_df)),
        "unique_users": int(filtered_df["username"].nunique()),
        "unique_anime": int(filtered_df["anime_id"].nunique()),
        "derived_user_features": int(len(feature_df)),
        "training_config": {
            "min_interactions_per_user": config.min_interactions_per_user,
            "min_interactions_per_item": config.min_interactions_per_item,
            "test_set_fraction": config.test_set_fraction,
            "n_components": config.n_components,
            "learning_rate": config.learning_rate,
            "loss_function": config.loss_function,
            "epochs": config.epochs,
            "threads": config.threads,
            "random_seed": config.random_seed,
        },
        "metrics": metrics,
    }
    write_json(summary, paths.training_report)
    console.success(f"Wrote training report to {paths.training_report}")
    return summary