#!/usr/bin/env python3
"""Reproduce PAVE's quarterly look-through valuation and earnings growth.

Data sources
------------
1. SEC Form N-PORT-P filings: quarter-end holdings, quantities and values.
2. SEC Company Facts API: point-in-time EPS and net income disclosed by each
   constituent no later than the holdings date.

The script writes ``backtest_result.json`` next to itself.  It deliberately
uses only public SEC endpoints and needs no API key.
"""

from __future__ import annotations

import json
import math
import os
import re
import time
import unicodedata
from collections import defaultdict
from datetime import date, datetime, timedelta
from difflib import SequenceMatcher
from pathlib import Path
from typing import Any

import requests
from lxml import etree


SEC_CIK = "1432353"
USER_AGENT = os.environ.get(
    "SEC_USER_AGENT", "chaoe.net PAVE research admin@chaoe.net"
)
CACHE_DIR = Path(os.environ.get("PAVE_CACHE", "/tmp/pave-sec-cache"))
OUT = Path(__file__).with_name("backtest_result.json")

# One PAVE filing per fiscal quarter.  The fund reports at Feb/May/Aug/Nov end.
# The accession list is intentionally explicit so the historical sample cannot
# silently change when SEC search ranking changes.
ACCESSIONS = [
    "0001752724-22-014707", "0001752724-22-096452",
    "0001752724-22-170226", "0001752724-22-238992",
    "0001752724-23-016590", "0001752724-23-092609",
    "0001752724-23-163587", "0001752724-23-242772",
    "0001752724-24-017800", "0001752724-24-094137",
    "0001752724-24-169552", "0001752724-24-244803",
    "0001752724-25-016362", "0001752724-25-096566",
    "0001752724-25-183074", "0002048251-25-001696",
    "0002048251-26-000833", "0002048251-26-003669",
]

# SEC's current ticker file omits some old names and legal-name variants.
# Values are current SEC tickers; delisted names are left unmatched rather than
# guessed, and their weight is disclosed as a coverage gap in the output.
ALIASES = {
    "ACUITY BRANDS": "AYI", "ACUITY": "AYI",
    "ALLEGHENY TECHNOLOGIES": "ATI", "AMRIZE AG": "AMRZ",
    "CRANE HOLDINGS": "CR", "EATON PUBLIC": "ETN",
    "H AND E EQUIPMENT SERVICES": "HEES", "HAYNES INTERNATIONAL": "HAYN",
    "IDEX": "IEX", "JACOBS ENGINEERING GROUP": "J",
    "NORTHWEST PIPE": "NWPX", "NOW": "DNOW",
    "PENTAIR PUBLIC": "PNR", "RELIANCE STEEL AND ALUMINUM": "RS",
    "SEMPRA ENERGY": "SRE", "SPX": "SPXC",
    "STERLING CONSTRUCTION": "STRL", "SUMMIT MATERIALS": "SUM",
    "TIMKENSTEEL": "TMST", "TOPBUILD": "BLD",
    "UNITED STATES STEEL": "X", "WESTLAKE CHEMICAL": "WLK",
    "ZURN WATER SOLUTIONS": "ZWS",
}

session = requests.Session()
session.headers.update({"User-Agent": USER_AGENT, "Accept-Encoding": "gzip, deflate"})


def cached_json(url: str, key: str) -> dict[str, Any]:
    CACHE_DIR.mkdir(parents=True, exist_ok=True)
    path = CACHE_DIR / f"{key}.json"
    if path.exists():
        return json.loads(path.read_text())
    for attempt in range(5):
        response = session.get(url, timeout=45)
        if response.status_code == 200:
            path.write_text(response.text)
            time.sleep(0.11)  # stay below SEC's fair-access ceiling
            return response.json()
        time.sleep(1.5 * (attempt + 1))
    response.raise_for_status()
    raise RuntimeError(url)


def cached_bytes(url: str, key: str) -> bytes:
    CACHE_DIR.mkdir(parents=True, exist_ok=True)
    path = CACHE_DIR / f"{key}.xml"
    if path.exists():
        return path.read_bytes()
    response = session.get(url, timeout=45)
    response.raise_for_status()
    path.write_bytes(response.content)
    time.sleep(0.11)
    return response.content


def normalize_name(value: str) -> str:
    value = unicodedata.normalize("NFKD", value).upper().replace("&", " AND ")
    value = re.sub(r"[^A-Z0-9 ]", " ", value)
    stop = {
        "INC", "INCORPORATED", "CORP", "CORPORATION", "CO", "COMPANY",
        "LTD", "LIMITED", "PLC", "LP", "LLC", "THE", "COMMON", "STOCK",
        "CLASS", "CL", "COM", "NEW", "DE", "DEL",
    }
    return " ".join(word for word in value.split() if word not in stop)


def parse_holdings(accession: str) -> dict[str, Any]:
    compact = accession.replace("-", "")
    url = f"https://www.sec.gov/Archives/edgar/data/{SEC_CIK}/{compact}/primary_doc.xml"
    root = etree.fromstring(cached_bytes(url, f"nport-{compact}"))
    one = lambda tag: next(
        (node.text or "").strip()
        for node in root.iter() if node.tag.split("}")[-1] == tag
    )
    if one("seriesName") != "Global X U.S. Infrastructure Development ETF":
        raise ValueError(f"{accession} is not PAVE")
    holdings = []
    for node in root.iter():
        if node.tag.split("}")[-1] != "invstOrSec":
            continue
        fields = {
            child.tag.split("}")[-1]: (child.text or "").strip()
            for child in node.iter()
        }
        if (
            fields.get("assetCat") == "EC"
            and fields.get("units") == "NS"
            and fields.get("invCountry") == "US"
            and fields.get("payoffProfile") == "Long"
        ):
            holdings.append({
                "name": fields["name"], "cusip": fields.get("cusip", ""),
                "shares": float(fields["balance"]),
                "value": float(fields["valUSD"]),
                "weight": float(fields["pctVal"]) / 100.0,
            })
    return {
        "date": one("repPdDate"), "net_assets": float(one("netAssets")),
        "accession": accession, "holdings": holdings,
    }


def ticker_master() -> tuple[dict[str, dict[str, Any]], list[dict[str, Any]]]:
    raw = cached_json("https://www.sec.gov/files/company_tickers.json", "company-tickers")
    rows = []
    by_ticker = {}
    for item in raw.values():
        row = {
            "ticker": item["ticker"].upper(), "cik": int(item["cik_str"]),
            "title": item["title"], "norm": normalize_name(item["title"]),
        }
        rows.append(row)
        by_ticker[row["ticker"]] = row
    return by_ticker, rows


def match_company(name: str, by_ticker: dict[str, Any], rows: list[dict[str, Any]]):
    norm = normalize_name(name)
    alias = ALIASES.get(norm)
    if alias and alias in by_ticker:
        return by_ticker[alias], 1.0, "alias"
    exact = [row for row in rows if row["norm"] == norm]
    if exact:
        return exact[0], 1.0, "exact"
    best = max(rows, key=lambda row: SequenceMatcher(None, norm, row["norm"]).ratio())
    score = SequenceMatcher(None, norm, best["norm"]).ratio()
    # A high threshold is preferable to silently attaching another issuer's EPS.
    if score >= 0.88:
        return best, score, "fuzzy"
    return None, score, "unmatched"


def parse_day(value: str) -> date:
    return datetime.strptime(value, "%Y-%m-%d").date()


def pick_unit(fact: dict[str, Any], preferred: tuple[str, ...]):
    for unit in preferred:
        if unit in fact.get("units", {}):
            return fact["units"][unit]
    return []


def select_records(companyfacts: dict[str, Any], concepts: tuple[str, ...], units):
    gaap = companyfacts.get("facts", {}).get("us-gaap", {})
    for concept in concepts:
        if concept in gaap:
            records = pick_unit(gaap[concept], units)
            if records:
                return records, concept
    return [], None


def latest_value(records: list[dict[str, Any]], cutoff: date) -> float | None:
    """Point-in-time TTM flow using FY + current YTD - comparable prior YTD."""
    clean = []
    for row in records:
        if row.get("form") not in {"10-Q", "10-K", "20-F", "40-F"}:
            continue
        try:
            start, end, filed = map(parse_day, (row["start"], row["end"], row["filed"]))
        except (KeyError, ValueError):
            continue
        if filed <= cutoff and end <= cutoff + timedelta(days=7):
            clean.append({**row, "_start": start, "_end": end, "_filed": filed,
                          "_days": (end - start).days})
    if not clean:
        return None

    # Retain the latest value known by the cutoff for each reported period.
    periods = {}
    for row in clean:
        key = (row["_start"], row["_end"])
        if key not in periods or row["_filed"] > periods[key]["_filed"]:
            periods[key] = row
    clean = list(periods.values())

    annuals = [row for row in clean if 300 <= row["_days"] <= 400]
    if not annuals:
        return None
    annual = max(annuals, key=lambda row: (row["_end"], row["_filed"]))

    ytds = [
        row for row in clean
        if row["_end"] > annual["_end"] and 55 <= row["_days"] <= 300
    ]
    if not ytds:
        return float(annual["val"])
    current = max(ytds, key=lambda row: (row["_end"], row["_days"], row["_filed"]))
    prior = [
        row for row in clean
        if 330 <= (current["_end"] - row["_end"]).days <= 400
        and abs(current["_days"] - row["_days"]) <= 25
    ]
    if not prior:
        return float(annual["val"])
    comparable = max(prior, key=lambda row: (row["_end"], row["_filed"]))
    return float(annual["val"]) + float(current["val"]) - float(comparable["val"])


def weighted_median(values: list[tuple[float, float]]) -> float | None:
    values = sorted((value, weight) for value, weight in values if math.isfinite(value))
    if not values:
        return None
    total = sum(weight for _, weight in values)
    running = 0.0
    for value, weight in values:
        running += weight
        if running >= total / 2:
            return value
    return values[-1][0]


def quarter_label(day: str) -> str:
    dt = parse_day(day)
    return f"{dt.year}Q{(dt.month - 1) // 3 + 1}"


def main() -> None:
    by_ticker, companies = ticker_master()
    snapshots = [parse_holdings(accession) for accession in ACCESSIONS]

    matches = {}
    for snapshot in snapshots:
        for holding in snapshot["holdings"]:
            name = holding["name"]
            if name not in matches:
                matches[name] = match_company(name, by_ticker, companies)

    ciks = sorted({match[0]["cik"] for match in matches.values() if match[0]})
    facts = {}
    for index, cik in enumerate(ciks, 1):
        facts[cik] = cached_json(
            f"https://data.sec.gov/api/xbrl/companyfacts/CIK{cik:010d}.json",
            f"companyfacts-{cik:010d}",
        )
        if index % 20 == 0:
            print(f"downloaded company facts: {index}/{len(ciks)}")

    output = []
    detail_latest = []
    unmatched_weights = defaultdict(float)
    for snapshot in snapshots:
        cutoff = parse_day(snapshot["date"])
        prior_cutoff = cutoff - timedelta(days=365)
        pe_value = pe_earnings = pe_weight = 0.0
        growth_values = []
        growth_weight = 0.0
        matched_weight = 0.0
        rows = []
        for holding in snapshot["holdings"]:
            company, score, method = matches[holding["name"]]
            row = {**holding, "match_score": round(score, 4), "match_method": method}
            if not company:
                unmatched_weights[holding["name"]] += holding["weight"]
                rows.append(row)
                continue
            matched_weight += holding["weight"]
            row.update({"ticker": company["ticker"], "cik": company["cik"]})
            companyfacts = facts[company["cik"]]
            eps_records, eps_concept = select_records(
                companyfacts,
                ("EarningsPerShareDiluted", "EarningsPerShareDilutedIncludingExtraordinaryItems",
                 "EarningsPerShareBasicAndDiluted", "EarningsPerShareBasic"),
                ("USD/shares", "USD / shares"),
            )
            ni_records, ni_concept = select_records(
                companyfacts, ("NetIncomeLoss", "ProfitLoss", "NetIncomeLossAvailableToCommonStockholdersBasic"),
                ("USD",),
            )
            eps = latest_value(eps_records, cutoff)
            net_income = latest_value(ni_records, cutoff)
            prior_income = latest_value(ni_records, prior_cutoff)
            row.update({"eps_ttm": eps, "eps_concept": eps_concept,
                        "net_income_ttm": net_income, "net_income_concept": ni_concept})
            if eps is not None and eps > 0:
                implied = holding["shares"] * eps
                pe_value += holding["value"]
                pe_earnings += implied
                pe_weight += holding["weight"]
                row["pe"] = holding["value"] / implied if implied > 0 else None
            if net_income is not None and prior_income is not None and net_income > 0 and prior_income > 0:
                growth = net_income / prior_income - 1.0
                growth_values.append((growth, holding["weight"]))
                growth_weight += holding["weight"]
                row["profit_growth"] = growth
            rows.append(row)

        winsorized = [
            (max(-1.0, min(2.0, value)), weight) for value, weight in growth_values
        ]
        growth_mean = (
            sum(value * weight for value, weight in winsorized) / sum(weight for _, weight in winsorized)
            if winsorized else None
        )
        result = {
            "quarter": quarter_label(snapshot["date"]), "date": snapshot["date"],
            "accession": snapshot["accession"], "net_assets_usd": snapshot["net_assets"],
            "equity_count": len(snapshot["holdings"]),
            "matched_weight": matched_weight, "pe_coverage": pe_weight,
            "growth_coverage": growth_weight,
            "pe_ttm": pe_value / pe_earnings if pe_earnings > 0 else None,
            "earnings_yield": pe_earnings / pe_value if pe_value > 0 else None,
            "profit_growth_weighted": growth_mean,
            "profit_growth_median": weighted_median(growth_values),
        }
        output.append(result)
        if snapshot is snapshots[-1]:
            detail_latest = sorted(rows, key=lambda row: row["weight"], reverse=True)

    result = {
        "generated_at": datetime.now().astimezone().isoformat(timespec="seconds"),
        "methodology": {
            "holdings": "SEC N-PORT-P fiscal quarter-end holdings (Feb/May/Aug/Nov)",
            "valuation": "Harmonic aggregate trailing P/E; positive-EPS covered holdings only",
            "growth": "Holding-weighted mean (constituent TTM net-income YoY winsorized to -100%/+200%) and weighted median",
            "point_in_time": "Only financial filings published on or before each holdings date",
        },
        "quarters": output,
        "latest_holdings": detail_latest,
        "unmatched": [
            {"name": name, "cumulative_weight": weight}
            for name, weight in sorted(unmatched_weights.items(), key=lambda item: -item[1])
        ],
    }
    OUT.write_text(json.dumps(result, ensure_ascii=False, indent=2) + "\n")
    print(f"wrote {OUT}")
    for row in output:
        print(row["quarter"], row["date"],
              f"PE={row['pe_ttm']:.2f}" if row["pe_ttm"] else "PE=n/a",
              f"growth={row['profit_growth_weighted']:.1%}" if row["profit_growth_weighted"] is not None else "growth=n/a",
              f"coverage={row['pe_coverage']:.1%}/{row['growth_coverage']:.1%}")


if __name__ == "__main__":
    main()
