# ============================================================
# 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

import argparse
import sys
import time
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

from src.embeddings.models import EmbeddingModelRegistry
from src.factors.anchors import AnchorRegistry
from src.factors.scoring import ScoringEngine
from src.factors.aggregation import FactorAggregator
from src.factors.longitudinal import LongitudinalBuilder
from src.factors.panels import PanelExporter
from src.factors.diagnostics import FactorDiagnostics
from src.storage.database import DatabaseManager
from src.utils.logging import setup_logging

logger = setup_logging("build_factors")

STEPS = ["all", "anchors", "scores", "aggregate", "longitudinal", "panels", "diagnostics"]


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Build thematic factor scores from SEC filing embeddings"
    )
    parser.add_argument(
        "--step",
        type=str,
        default="all",
        choices=STEPS,
        help="Pipeline step to run (default: all)",
    )
    parser.add_argument(
        "--batch-size",
        type=int,
        default=10_000,
        help="Batch size for scoring (default: 10000)",
    )
    parser.add_argument(
        "--skip-drift",
        action="store_true",
        default=False,
        help="Skip embedding drift computation (slower step)",
    )
    parser.add_argument(
        "--skip-novelty",
        action="store_true",
        default=False,
        help="Skip novelty score computation",
    )
    return parser.parse_args()


def run_step(name: str) -> None:
    logger.info("=" * 60)
    logger.info("STEP: %s", name.upper())
    logger.info("=" * 60)


def main() -> None:
    args = parse_args()
    t0 = time.time()

    db = DatabaseManager()
    db.initialize()
    registry = EmbeddingModelRegistry()

    anchor_reg = AnchorRegistry(registry)
    anchor_ids, anchor_matrix = anchor_reg.get_anchor_matrix()
    logger.info("Anchor IDs: %s", anchor_ids)
    logger.info("Anchor matrix shape: %s", anchor_matrix.shape)

    # ---- Step 1: Anchors ----
    if args.step in ("all", "anchors"):
        run_step("anchors")
        anchor_reg.save_to_duckdb(db)
        anchor_reg.save_to_parquet()
        logger.info("Anchors embedded and saved")

    # ---- Step 2: Scoring ----
    if args.step in ("all", "scores"):
        run_step("scoring")
        engine = ScoringEngine(db, batch_size=args.batch_size)
        scores_df = engine.compute_scores(anchor_ids, anchor_matrix)
        engine.save_chunk_scores(scores_df)
        del scores_df
        logger.info("Chunk scoring complete")

    # ---- Step 3: Aggregation ----
    if args.step in ("all", "aggregate"):
        run_step("aggregation")
        agg = FactorAggregator(db)

        filing_df = agg.aggregate_to_filing(anchor_ids)
        agg.save_filing_scores(filing_df)

        quarterly_df = agg.aggregate_to_quarterly(anchor_ids)
        agg.save_quarterly_panel(quarterly_df)

        annual_df = agg.aggregate_to_annual(anchor_ids)
        agg.save_annual_panel(annual_df)
        logger.info("Aggregation complete")

    # ---- Step 4: Longitudinal ----
    drift_df = None
    novelty_df = None
    if args.step in ("all", "longitudinal"):
        run_step("longitudinal")
        longi = LongitudinalBuilder(db)

        if not args.skip_drift:
            drift_df = longi.compute_embedding_drift()
            logger.info("Embedding drift computed")
        else:
            logger.info("Skipping embedding drift (--skip-drift)")

        if not args.skip_novelty:
            novelty_df = longi.compute_novelty()
            logger.info("Novelty scores computed")
        else:
            logger.info("Skipping novelty (--skip-novelty)")

    # ---- Step 5: Panels ----
    if args.step in ("all", "panels"):
        run_step("panels")
        exporter = PanelExporter(db)

        annual = exporter.build_annual_panel(anchor_ids, drift_df, novelty_df)
        quarterly = exporter.build_quarterly_panel(anchor_ids)
        exporter.export_all(annual, quarterly)
        logger.info("Panel export complete")

    # ---- Step 6: Diagnostics ----
    if args.step in ("all", "diagnostics"):
        run_step("diagnostics")
        try:
            annual = db.fetch_df("SELECT * FROM annual_final")
        except Exception:
            annual = db.fetch_df("SELECT * FROM annual_panel")
            companies = db.fetch_df("SELECT ticker, sector, sub_industry FROM companies")
            annual = annual.merge(companies, on="ticker", how="left")

        diag = FactorDiagnostics(db)
        diag.save_diagnostics(annual, anchor_ids)
        logger.info("Diagnostics complete")

    elapsed = time.time() - t0
    logger.info("=" * 60)
    logger.info("PIPELINE COMPLETE in %.1f seconds (%.1f minutes)", elapsed, elapsed / 60)
    logger.info("=" * 60)

    stats = {}
    for table in ["anchor_embeddings", "chunk_scores", "filing_scores", "quarterly_panel", "annual_panel"]:
        try:
            count = db.execute(f"SELECT COUNT(*) FROM {table}").fetchone()[0]
            stats[table] = count
        except Exception:
            stats[table] = "N/A"
    logger.info("Final table counts:")
    for table, count in stats.items():
        logger.info("  %-25s %s rows", table, f"{count:,}" if isinstance(count, int) else count)

    db.close()


if __name__ == "__main__":
    main()