# ============================================================
# Project: EDGAR Research Database
# Description: SEC EDGAR Research Infrastructure
#
# Authors:
# - Manel Kammoun — manel.kammoun@uqo.ca
# - Djerrycharli Tandjambianda — djerrycharli.tandjambianda@uqo.ca
# - Simon-Pierre Boucher — simon-pierre.boucher@uqo.ca
#
# Institution:
# Universite du Quebec en Outaouais (UQO)
#
# License:
# Research / Academic Use
# ============================================================

from __future__ import annotations

from typing import TYPE_CHECKING

import numpy as np
import pandas as pd
import pyarrow.dataset as ds
from tqdm import tqdm

from src.utils.config import get_config
from src.utils.constants import DATA_DIRS
from src.utils.logging import get_logger

if TYPE_CHECKING:
    from src.storage.database import DatabaseManager

logger = get_logger("factors.longitudinal")


def _firm_year_centroids(chunk_to_firmyear: dict) -> dict[tuple[str, int], np.ndarray]:
    """Mean (then unit-normalised) embedding per (ticker, year), streamed from the parquet store."""
    from src.factors.scoring import _vectors_from_column

    keys = sorted(set(chunk_to_firmyear.values()), key=lambda k: (str(k[0]), int(k[1]) if k[1] is not None else -1))
    key_index = {k: i for i, k in enumerate(keys)}
    chunk_index = {cid: key_index[k] for cid, k in chunk_to_firmyear.items()}

    sums: np.ndarray | None = None
    counts = np.zeros(len(keys), dtype=np.int64)

    emb_dir = DATA_DIRS["parquet"] / "embeddings"
    dataset = ds.dataset(str(emb_dir), format="parquet", partitioning="hive")
    scanner = dataset.scanner(columns=["source_id", "vector"], batch_size=20_000)
    for batch in tqdm(scanner.to_batches(), desc="Accumulating firm-year centroids", unit="batch"):
        if batch.num_rows == 0:
            continue
        ids = batch.column("source_id").to_pylist()
        idx = np.fromiter((chunk_index.get(i, -1) for i in ids), dtype=np.int64, count=len(ids))
        mask = idx >= 0
        if not mask.any():
            continue
        vectors = _vectors_from_column(batch.column("vector"))[mask].astype(np.float64)
        if sums is None:
            sums = np.zeros((len(keys), vectors.shape[1]), dtype=np.float64)
        np.add.at(sums, idx[mask], vectors)
        np.add.at(counts, idx[mask], 1)

    centroids: dict[tuple[str, int], np.ndarray] = {}
    if sums is None:
        return centroids
    for k, i in key_index.items():
        if counts[i] == 0:
            continue
        c = (sums[i] / counts[i]).astype(np.float32)
        n = np.linalg.norm(c)
        centroids[k] = c / n if n > 0 else c
    return centroids


class LongitudinalBuilder:

    def __init__(self, db: DatabaseManager) -> None:
        self.db = db
        cfg = get_config("factors").get("factors", {}).get("longitudinal", {})
        self._rolling_window = cfg.get("rolling_window", 3)
        self._min_years_drift = cfg.get("min_years_for_drift", 2)

    def add_yoy_changes(
        self,
        panel: pd.DataFrame,
        anchor_ids: list[str],
    ) -> pd.DataFrame:
        panel = panel.sort_values(["ticker", "fiscal_year"]).copy()
        for aid in anchor_ids:
            col = f"{aid}_mean"
            if col not in panel.columns:
                continue
            panel[f"{aid}_delta"] = panel.groupby("ticker")[col].diff(1)
            panel[f"{aid}_pct_change"] = panel.groupby("ticker")[col].pct_change(1)
        logger.info("Added YoY changes for %d anchors", len(anchor_ids))
        return panel

    def add_acceleration(
        self,
        panel: pd.DataFrame,
        anchor_ids: list[str],
    ) -> pd.DataFrame:
        panel = panel.sort_values(["ticker", "fiscal_year"]).copy()
        for aid in anchor_ids:
            delta_col = f"{aid}_delta"
            if delta_col not in panel.columns:
                continue
            panel[f"{aid}_accel"] = panel.groupby("ticker")[delta_col].diff(1)
        logger.info("Added acceleration for %d anchors", len(anchor_ids))
        return panel

    def add_rolling_stats(
        self,
        panel: pd.DataFrame,
        anchor_ids: list[str],
        window: int | None = None,
    ) -> pd.DataFrame:
        w = window or self._rolling_window
        panel = panel.sort_values(["ticker", "fiscal_year"]).copy()
        for aid in anchor_ids:
            col = f"{aid}_mean"
            if col not in panel.columns:
                continue
            grp = panel.groupby("ticker")[col]
            panel[f"{aid}_roll{w}_mean"] = grp.transform(
                lambda x: x.rolling(w, min_periods=2).mean()
            )
            panel[f"{aid}_roll{w}_std"] = grp.transform(
                lambda x: x.rolling(w, min_periods=2).std()
            )
        logger.info("Added rolling stats (window=%d) for %d anchors", w, len(anchor_ids))
        return panel

    def compute_embedding_drift(self) -> pd.DataFrame:
        logger.info("Computing firm-year embedding centroids for drift...")

        chunk_years = self.db.fetch_df("""
            SELECT
                c.chunk_id,
                f.ticker,
                EXTRACT(YEAR FROM f.filing_date)::INTEGER AS fiscal_year
            FROM filing_chunks c
            JOIN filings f ON c.accession_number = f.accession_number
        """)

        chunk_to_firmyear = dict(
            zip(chunk_years["chunk_id"], zip(chunk_years["ticker"], chunk_years["fiscal_year"]))
        )

        centroids = _firm_year_centroids(chunk_to_firmyear)
        del chunk_to_firmyear

        tickers = sorted(set(k[0] for k in centroids))
        drift_records = []
        for ticker in tickers:
            years = sorted(y for t, y in centroids if t == ticker)
            for i in range(1, len(years)):
                y_prev, y_curr = years[i - 1], years[i]
                if y_curr - y_prev > 2:
                    continue
                c_prev = centroids[(ticker, y_prev)]
                c_curr = centroids[(ticker, y_curr)]
                sim = float(np.dot(c_prev, c_curr))
                drift_records.append({
                    "ticker": ticker,
                    "fiscal_year": y_curr,
                    "embedding_drift": 1.0 - sim,
                })

        drift_df = pd.DataFrame(drift_records)
        logger.info("Computed embedding drift for %d firm-year observations", len(drift_df))
        return drift_df

    def compute_novelty(self) -> pd.DataFrame:
        logger.info("Computing novelty scores (distance from sector centroid)...")

        companies = self.db.fetch_df("SELECT cik, ticker, sector FROM companies")
        ticker_to_sector = dict(zip(companies["ticker"], companies["sector"]))

        annual = self.db.fetch_df("SELECT ticker, fiscal_year FROM annual_panel")

        chunk_years = self.db.fetch_df("""
            SELECT c.chunk_id, f.ticker,
                   EXTRACT(YEAR FROM f.filing_date)::INTEGER AS fiscal_year
            FROM filing_chunks c
            JOIN filings f ON c.accession_number = f.accession_number
        """)
        chunk_to_firmyear = dict(
            zip(chunk_years["chunk_id"], zip(chunk_years["ticker"], chunk_years["fiscal_year"]))
        )

        firm_centroids = _firm_year_centroids(chunk_to_firmyear)
        del chunk_to_firmyear

        sector_year_sums: dict[tuple[str, int], np.ndarray] = {}
        sector_year_counts: dict[tuple[str, int], int] = {}
        for (ticker, year), centroid in firm_centroids.items():
            sector = ticker_to_sector.get(ticker, "Unknown")
            key = (sector, year)
            if key not in sector_year_sums:
                sector_year_sums[key] = np.zeros_like(centroid, dtype=np.float64)
                sector_year_counts[key] = 0
            sector_year_sums[key] += centroid.astype(np.float64)
            sector_year_counts[key] += 1

        sector_centroids: dict[tuple[str, int], np.ndarray] = {}
        for key, s in sector_year_sums.items():
            centroid = (s / sector_year_counts[key]).astype(np.float32)
            norm = np.linalg.norm(centroid)
            if norm > 0:
                centroid = centroid / norm
            sector_centroids[key] = centroid

        novelty_records = []
        for (ticker, year), firm_c in firm_centroids.items():
            sector = ticker_to_sector.get(ticker, "Unknown")
            sec_c = sector_centroids.get((sector, year))
            if sec_c is None:
                continue
            sim = float(np.dot(firm_c, sec_c))
            novelty_records.append({
                "ticker": ticker,
                "fiscal_year": year,
                "novelty_score": 1.0 - sim,
            })

        novelty_df = pd.DataFrame(novelty_records)
        logger.info("Computed novelty for %d firm-year observations", len(novelty_df))
        return novelty_df