from __future__ import annotations

from collections import defaultdict
from statistics import fmean, median
from typing import Any

from backend.features.screener.catalog import factor_catalog
from backend.features.screener.factor_math import mean, number, pearson, percentile_map, rounded


def finalize_factor_rows(
    rows: list[dict[str, Any]], dataset_ready: dict[str, bool]
) -> list[dict[str, Any]]:
    market_returns = [number(row.get("return_5d")) for row in rows]
    valid_market = [value for value in market_returns if value is not None]
    market_mean = fmean(valid_market) if valid_market else None
    for row in rows:
        stock_return = number(row.get("return_5d"))
        row["relative_strength"] = (
            rounded(stock_return - market_mean, 2)
            if stock_return is not None and market_mean is not None
            else None
        )
    _rank(rows, "return_5d", "return_5d_rank", "desc")
    _rank(rows, "momentum_60_5", "momentum_60_5_rank", "desc")
    _sector_factors(rows, dataset_ready)
    _composite_factors(rows)
    _style_factors(rows)
    _market_height(rows)
    known = factor_catalog()["factors"]
    for row in rows:
        for field in known:
            row.setdefault(field, None)
    return rows


def _sector_factors(rows: list[dict[str, Any]], dataset_ready: dict[str, bool]) -> None:
    fields = (
        "sector_strength",
        "sector_return_5d",
        "sector_return_20d",
        "sector_momentum_rank",
        "sector_stock_momentum_rank",
        "sector_net_flow_5d_million",
        "sector_flow_rank",
        "sector_prosperity_rank",
        "sector_trend_rank",
        "sector_crowding_rank",
        "sector_composite_score",
        "sector_limit_count",
        "sector_up_count",
        "sector_breadth_ma20",
    )
    if not dataset_ready.get("industry"):
        for row in rows:
            row.update(dict.fromkeys(fields))
        return
    sectors: dict[str, list[dict[str, Any]]] = defaultdict(list)
    for row in rows:
        if row.get("sector"):
            sectors[str(row["sector"])].append(row)
    market_amount = sum(float(number(row.get("amount_billion")) or 0) for row in rows)
    metrics = []
    for name, members in sectors.items():
        returns_5 = [number(row.get("return_5d")) for row in members]
        returns_20 = [number(row.get("return_20d")) for row in members]
        average_5 = mean(returns_5)
        average_20 = mean(returns_20)
        flows = [number(row.get("net_flow_5d_million")) for row in members]
        sector_flow = (
            sum(float(value) for value in flows)
            if all(value is not None for value in flows)
            else None
        )
        limit_values = [row.get("is_limit_up_today") for row in members]
        limit_count = (
            sum(bool(value) for value in limit_values)
            if all(value is not None for value in limit_values)
            else None
        )
        changes = [number(row.get("pct_chg")) for row in members]
        up_count = (
            sum(float(value) >= 5 for value in changes if value is not None)
            if all(value is not None for value in changes)
            else None
        )
        above = [row.get("above_ma20") for row in members]
        breadth = (
            sum(bool(value) for value in above) / len(above) * 100
            if above and all(value is not None for value in above)
            else None
        )
        growth = [
            mean([number(row.get("revenue_yoy")), number(row.get("netprofit_yoy"))])
            for row in members
        ]
        valid_growth = [value for value in growth if value is not None]
        prosperity = median(valid_growth) if valid_growth else None
        turnovers = [number(row.get("turnover_rate")) for row in members]
        average_turnover = mean(turnovers)
        amount_share = (
            sum(float(number(row.get("amount_billion")) or 0) for row in members)
            / market_amount
            * 100
            if market_amount
            else None
        )
        crowding = (
            average_turnover + amount_share
            if average_turnover is not None and amount_share is not None
            else None
        )
        trend = (
            average_20 + breadth / 10 if average_20 is not None and breadth is not None else None
        )
        strength = (
            min(
                100,
                max(
                    0,
                    50 + average_5 * 4 + (limit_count or 0) * 3 + (up_count or 0) * 0.6,
                ),
            )
            if average_5 is not None
            else None
        )
        metrics.append(
            {
                "identifier": name,
                "sector": name,
                "return_20": average_20,
                "flow": sector_flow,
                "prosperity": prosperity,
                "trend": trend,
                "crowding": crowding,
            }
        )
        stock_ranks = percentile_map(members, "return_20d", "desc")
        for row in members:
            row.update(
                {
                    "sector_strength": rounded(strength, 1),
                    "sector_return_5d": rounded(average_5, 2),
                    "sector_return_20d": rounded(average_20, 2),
                    "sector_stock_momentum_rank": rounded(
                        stock_ranks.get(str(row["identifier"])), 4
                    ),
                    "sector_net_flow_5d_million": rounded(sector_flow, 2),
                    "sector_limit_count": limit_count,
                    "sector_up_count": up_count,
                    "sector_breadth_ma20": rounded(breadth, 1),
                }
            )
    rank_specs = {
        "sector_momentum_rank": ("return_20", "desc"),
        "sector_flow_rank": ("flow", "desc"),
        "sector_prosperity_rank": ("prosperity", "desc"),
        "sector_trend_rank": ("trend", "desc"),
        "sector_crowding_rank": ("crowding", "desc"),
    }
    maps = {
        output: percentile_map(metrics, source, direction)
        for output, (source, direction) in rank_specs.items()
    }
    for name, members in sectors.items():
        values = {field: mapping.get(name) for field, mapping in maps.items()}
        composite = (
            values["sector_prosperity_rank"] * 0.4
            + values["sector_trend_rank"] * 0.3
            + (1 - values["sector_crowding_rank"]) * 0.3
            if all(value is not None for value in values.values())
            else None
        )
        for row in members:
            row.update({field: rounded(value, 4) for field, value in values.items()})
            row["sector_composite_score"] = rounded(composite, 4)
    for row in rows:
        if not row.get("sector"):
            row.update(dict.fromkeys(fields))


def _composite_factors(rows: list[dict[str, Any]]) -> None:
    specs = {
        "factor_value_score": (("pe_ttm", "asc"), ("pb", "asc"), ("dividend_yield_ttm", "desc")),
        "factor_growth_score": (("revenue_yoy", "desc"), ("netprofit_yoy", "desc")),
        "factor_quality_score": (("roe", "desc"), ("roic", "desc"), ("gross_margin", "desc")),
        "factor_momentum_score": (("momentum_60_5", "desc"), ("relative_strength", "desc")),
        "factor_sentiment_score": (("turnover_rate", "desc"), ("volume_ratio_5d", "desc")),
    }
    for output, factor_specs in specs.items():
        maps = [percentile_map(rows, field, direction) for field, direction in factor_specs]
        for row in rows:
            values = [mapping.get(str(row["identifier"])) for mapping in maps]
            row[output] = rounded(mean(values), 4)
    future_rank = percentile_map(rows, "return_20d", "desc")
    weights = {}
    for output in specs:
        pairs = [(number(row.get(output)), future_rank.get(str(row["identifier"]))) for row in rows]
        valid = [(left, right) for left, right in pairs if left is not None and right is not None]
        correlation = pearson(
            [float(left) for left, _ in valid],
            [float(right) for _, right in valid],
        )
        weights[output] = max(0.05, correlation)
    for row in rows:
        available = [
            (number(row.get(field)), weight)
            for field, weight in weights.items()
            if number(row.get(field)) is not None
        ]
        row["multi_factor_composite"] = (
            rounded(
                sum(float(value) * weight for value, weight in available)
                / sum(weight for _, weight in available),
                4,
            )
            if available
            else None
        )


def _style_factors(rows: list[dict[str, Any]]) -> None:
    size = percentile_map(rows, "total_mv_billion", "desc")
    large = [row for row in rows if (size.get(str(row["identifier"])) or 0) >= 0.7]
    small = [row for row in rows if (size.get(str(row["identifier"])) or 1) <= 0.3]
    large_return = mean([number(row.get("return_20d")) for row in large])
    small_return = mean([number(row.get("return_20d")) for row in small])
    prefer_large = (
        large_return >= small_return
        if large_return is not None and small_return is not None
        else None
    )
    growth = [row for row in rows if (number(row.get("factor_growth_score")) or 0) >= 0.7]
    value = [row for row in rows if (number(row.get("factor_value_score")) or 0) >= 0.7]
    growth_return = mean([number(row.get("return_20d")) for row in growth])
    value_return = mean([number(row.get("return_20d")) for row in value])
    prefer_growth = (
        growth_return >= value_return
        if growth_return is not None and value_return is not None
        else None
    )
    for row in rows:
        size_rank = size.get(str(row["identifier"]))
        row["style_size_fit"] = (
            rounded(size_rank if prefer_large else 1 - size_rank, 4)
            if size_rank is not None and prefer_large is not None
            else None
        )
        row["style_growth_fit"] = (
            row.get("factor_growth_score")
            if prefer_growth
            else row.get("factor_value_score")
            if prefer_growth is not None
            else None
        )
        row["style_fit_score"] = rounded(
            mean([number(row.get("style_size_fit")), number(row.get("style_growth_fit"))]),
            4,
        )


def _market_height(rows: list[dict[str, Any]]) -> None:
    current = [int(row["limit_streak"]) for row in rows if row.get("limit_streak") is not None]
    previous = [
        int(row["previous_limit_streak"])
        for row in rows
        if row.get("previous_limit_streak") is not None
    ]
    current_height = max(current, default=0)
    previous_height = max(previous, default=0)
    for row in rows:
        streak = row.get("limit_streak")
        prior = row.get("previous_limit_streak")
        if streak is None or prior is None:
            row["is_market_height"] = None
            row["new_space_board"] = None
            continue
        is_height = current_height >= 2 and int(streak) == current_height
        row["is_market_height"] = is_height
        row["new_space_board"] = is_height and not (
            previous_height >= 2 and int(prior) == previous_height
        )


def _rank(rows: list[dict[str, Any]], source: str, target: str, direction: str) -> None:
    mapping = percentile_map(rows, source, direction)
    for row in rows:
        row[target] = rounded(mapping.get(str(row["identifier"])), 4)
