# ============================================================
# 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 pathlib import Path
from typing import TYPE_CHECKING

import numpy as np
import pandas as pd
import pyarrow as pa
import pyarrow.parquet as pq

from src.utils.config import get_config
from src.utils.constants import DATA_DIRS, PARQUET_COMPRESSION, PARQUET_COMPRESSION_LEVEL, PROJECT_ROOT
from src.utils.logging import get_logger

if TYPE_CHECKING:
    from src.storage.database import DatabaseManager

logger = get_logger("factors.panels")


class PanelExporter:

    def __init__(self, db: DatabaseManager) -> None:
        self.db = db
        cfg = get_config("factors").get("factors", {}).get("export", {})
        self._formats = cfg.get("formats", ["parquet", "csv"])
        self._output_dir = PROJECT_ROOT / cfg.get("output_dir", "outputs/factors")

    def build_annual_panel(
        self,
        anchor_ids: list[str],
        drift_df: pd.DataFrame | None = None,
        novelty_df: pd.DataFrame | None = None,
    ) -> pd.DataFrame:
        logger.info("Building comprehensive annual panel...")

        annual = self.db.fetch_df("SELECT * FROM annual_panel")

        companies = self.db.fetch_df(
            "SELECT cik, ticker, company_name, sector, sub_industry FROM companies"
        )
        annual = annual.merge(
            companies[["ticker", "company_name", "sector", "sub_industry"]],
            on="ticker",
            how="left",
        )

        from src.factors.longitudinal import LongitudinalBuilder
        longi = LongitudinalBuilder(self.db)
        annual = longi.add_yoy_changes(annual, anchor_ids)
        annual = longi.add_acceleration(annual, anchor_ids)
        annual = longi.add_rolling_stats(annual, anchor_ids)

        if drift_df is not None and not drift_df.empty:
            annual = annual.merge(drift_df, on=["ticker", "fiscal_year"], how="left")
        else:
            annual["embedding_drift"] = np.nan

        if novelty_df is not None and not novelty_df.empty:
            annual = annual.merge(novelty_df, on=["ticker", "fiscal_year"], how="left")
        else:
            annual["novelty_score"] = np.nan

        annual["log_total_chunks"] = np.log1p(annual["total_chunks"])

        id_cols = [
            "cik", "ticker", "company_name", "sector", "sub_industry", "fiscal_year",
            "filing_count", "total_chunks", "log_total_chunks",
        ]
        score_cols = sorted([c for c in annual.columns if any(
            c.startswith(f"{aid}_") for aid in anchor_ids
        )])
        feature_cols = ["embedding_drift", "novelty_score"]
        existing = [c for c in id_cols + score_cols + feature_cols if c in annual.columns]
        annual = annual[existing]

        annual = annual.sort_values(["ticker", "fiscal_year"]).reset_index(drop=True)
        logger.info(
            "Annual panel: %d rows x %d columns (%d firms, %d years)",
            len(annual), len(annual.columns),
            annual["ticker"].nunique(),
            annual["fiscal_year"].nunique(),
        )
        return annual

    def build_quarterly_panel(self, anchor_ids: list[str]) -> pd.DataFrame:
        quarterly = self.db.fetch_df("SELECT * FROM quarterly_panel")
        companies = self.db.fetch_df(
            "SELECT ticker, company_name, sector, sub_industry FROM companies"
        )
        quarterly = quarterly.merge(companies, on="ticker", how="left")
        quarterly = quarterly.sort_values(
            ["ticker", "fiscal_year", "fiscal_quarter"]
        ).reset_index(drop=True)
        logger.info("Quarterly panel: %d rows x %d columns", len(quarterly), len(quarterly.columns))
        return quarterly

    def save_final_panel(self, panel: pd.DataFrame, name: str) -> None:
        self.db.execute(f"DROP TABLE IF EXISTS {name}_final")
        self.db.conn.execute(f"CREATE TABLE {name}_final AS SELECT * FROM panel")
        logger.info("Saved %s_final (%d rows) to DuckDB", name, len(panel))

        out_dir = DATA_DIRS["parquet"] / f"{name}_final"
        out_dir.mkdir(parents=True, exist_ok=True)
        table = pa.Table.from_pandas(panel, preserve_index=False)
        pq.write_table(
            table,
            str(out_dir / f"{name}_final.parquet"),
            compression=PARQUET_COMPRESSION,
            compression_level=PARQUET_COMPRESSION_LEVEL,
        )

        self._output_dir.mkdir(parents=True, exist_ok=True)
        if "csv" in self._formats:
            csv_path = self._output_dir / f"{name}_panel.csv"
            panel.to_csv(csv_path, index=False)
            logger.info("Exported %s to CSV: %s", name, csv_path)

        if "parquet" in self._formats:
            pq_path = self._output_dir / f"{name}_panel.parquet"
            pq.write_table(
                table,
                str(pq_path),
                compression=PARQUET_COMPRESSION,
                compression_level=PARQUET_COMPRESSION_LEVEL,
            )
            logger.info("Exported %s to parquet: %s", name, pq_path)

    def export_all(
        self,
        annual: pd.DataFrame,
        quarterly: pd.DataFrame,
    ) -> None:
        self.save_final_panel(annual, "annual")
        self.save_final_panel(quarterly, "quarterly")
        logger.info("All panels exported")