from __future__ import annotations

from collections import defaultdict
from datetime import date
from typing import Any

from backend.features.screener.factor_math import number, ratio, rounded


def group(rows: Any, field: str) -> dict[str, list[dict[str, Any]]]:
    result: dict[str, list[dict[str, Any]]] = defaultdict(list)
    for row in rows:
        key = str(row.get(field) or "")
        if key:
            result[key].append(dict(row))
    return result


def latest_by_code(rows: Any, through: str) -> dict[str, dict[str, Any]]:
    result: dict[str, dict[str, Any]] = {}
    for row in rows:
        identifier = str(row.get("ts_code") or "")
        row_date = str(row.get("trade_date") or "")
        if (
            identifier
            and row_date <= through
            and (
                identifier not in result
                or row_date > str(result[identifier].get("trade_date") or "")
            )
        ):
            result[identifier] = dict(row)
    return result


def point_in_time(rows: Any, through: str) -> dict[str, dict[str, Any]]:
    result: dict[str, dict[str, Any]] = {}
    for row in rows:
        identifier = str(row.get("ts_code") or "")
        announced = str(row.get("ann_date") or "")
        if (
            identifier
            and announced
            and announced <= through
            and (
                identifier not in result
                or announced > str(result[identifier].get("ann_date") or "")
            )
        ):
            result[identifier] = dict(row)
    return result


def limit_events(rows: Any) -> dict[str, dict[str, str]]:
    result: dict[str, dict[str, str]] = defaultdict(dict)
    for row in rows:
        date_value = str(row.get("trade_date") or "")
        identifier = str(row.get("ts_code") or "")
        event = str(row.get("limit_type") or "")
        if date_value and identifier and event in {"U", "D", "Z"}:
            result[date_value][identifier] = event
    return result


def listed_days(value: Any, through: str) -> int:
    try:
        listed = str(value or "").replace("-", "")
        start = date(int(listed[:4]), int(listed[4:6]), int(listed[6:]))
        return (date.fromisoformat(through) - start).days
    except (ValueError, TypeError):
        return 0


def dividend_years(rows: list[dict[str, Any]], through: str) -> int:
    return len(
        {
            str(row.get("end_date") or "")[:4]
            for row in rows
            if str(row.get("ann_date") or row.get("ex_date") or "") <= through
            and (number(row.get("cash_div_tax")) or 0) > 0
        }
    )


def large_flow(row: dict[str, Any]) -> float | None:
    if not row:
        return None
    direct = number(row.get("large_net_amount"))
    if direct is not None:
        return rounded(direct / 100, 2)
    buys = [number(row.get(field)) for field in ("buy_lg_amount", "buy_elg_amount")]
    sells = [number(row.get(field)) for field in ("sell_lg_amount", "sell_elg_amount")]
    if any(value is None for value in buys + sells):
        return None
    return rounded(
        (sum(float(value) for value in buys) - sum(float(value) for value in sells)) / 100,
        2,
    )


def ending_streak(flags: list[bool], end: int | None = None) -> int:
    index = len(flags) - 1 if end is None else end
    count = 0
    while index >= 0 and flags[index]:
        count += 1
        index -= 1
    return count


def max_streak(flags: list[bool]) -> int:
    best = current = 0
    for flag in flags:
        current = current + 1 if flag else 0
        best = max(best, current)
    return best


def broken_metrics(bars: list[dict[str, Any]], up_flags: list[bool]) -> dict[str, Any]:
    if len(bars) < 3:
        return {"signal": None, "days": None, "recovered": None, "volume_ratio": None}
    last_limit = next((index for index in range(len(up_flags) - 2, -1, -1) if up_flags[index]), -1)
    if last_limit < 0:
        return {"signal": False, "days": None, "recovered": False, "volume_ratio": None}
    days = len(bars) - 1 - last_limit
    broken_high = number(bars[last_limit].get("high"))
    recovered = (
        number(bars[-1].get("close")) is not None
        and broken_high is not None
        and float(number(bars[-1]["close"]) or 0) > broken_high
    )
    volume_ratio = ratio(number(bars[-1].get("vol")), number(bars[last_limit].get("vol")))
    return {
        "signal": 1 <= days <= 3 and recovered and volume_ratio is not None and volume_ratio >= 1,
        "days": days,
        "recovered": recovered,
        "volume_ratio": rounded(volume_ratio, 3),
    }
