"""Portfolio construction: weighted blend of assets with rebalancing.

All inputs are date x symbol DataFrames of *adjusted* prices (which
already include distributions). Returns are pre-tax.

Supported schemes:
    rebalance=None      buy & hold (drift for the whole period)
    rebalance='1W'/'1ME'/'QE'/'YE' etc. (pandas offset alias)
                        rebalance to target weights on the first trading
                        day of each period
    cost_bps            one-way trading cost in basis points, applied on
                        the traded fraction (turnover) at each rebalance

The weight evolution between rebalances is the standard drift:
    w_{t+1} = w_t (1 + r_t) / (1 + R_t)
"""

from __future__ import annotations

import math
from dataclasses import dataclass, field

import numpy as np
import pandas as pd


@dataclass
class PortfolioResult:
    equity: pd.Series                 # gross value, starts at 1.0
    returns: pd.Series                # daily portfolio return
    allocation: pd.DataFrame          # drifted weights (date x asset)
    turnover: pd.Series               # one-way turnover at each rebalance (else 0)
    costs: pd.Series                  # cost paid at each rebalance (else 0)
    rebalance_dates: list = field(default_factory=list)


def _rebalance_dates(index: pd.DatetimeIndex, freq: str) -> set:
    """First trading day of each period, including the first day."""
    periods = index.to_series().groupby(pd.Grouper(freq=freq)).first()
    return set(periods.index.tolist())


def portfolio_returns(adj: pd.DataFrame, weights: dict[str, float],
                      rebalance: str | None = None,
                      cost_bps: float = 0.0,
                      start: str | None = None, end: str | None = None
                      ) -> PortfolioResult:
    """Simulate a (drifting) weighted portfolio on an adjusted price panel.

    weights: {symbol: raw weight}; normalized internally.
    """
    p = adj[list(weights)].copy()
    if start or end:
        p = p.loc[start:end]
    p = p.dropna()
    if len(p) < 2:
        raise ValueError("no overlapping data for the given symbols/period")

    w0 = np.array([weights[s] for s in p.columns], dtype=float)
    w0 = w0 / w0.sum()

    rets = p.pct_change().fillna(0.0).values
    rebal_set = _rebalance_dates(p.index, rebalance) if rebalance else set()
    cost = cost_bps / 1e4

    n_days = len(p)
    w = w0.copy()
    eq = np.empty(n_days)
    alloc = np.empty((n_days, len(p.columns)))
    turns = np.zeros(n_days)
    costs = np.zeros(n_days)
    eq[0] = 1.0
    alloc[0] = w
    # initial buy: cost on 100% invested
    eq[0] *= (1.0 - cost * 1.0)
    costs[0] = cost

    for i in range(1, n_days):
        R = float(np.dot(w, rets[i]))
        w_drifted = w * (1.0 + rets[i]) / (1.0 + R)

        if p.index[i] in rebal_set:
            target = w0
            turnover = float(np.abs(target - w_drifted).sum() / 2.0)
            c = cost * turnover
            w_drifted = target
            eq[i] = eq[i - 1] * (1.0 + R) * (1.0 - c)
            turns[i] = turnover
            costs[i] = c
        else:
            eq[i] = eq[i - 1] * (1.0 + R)

        w = w_drifted
        alloc[i] = w

    idx = p.index
    return PortfolioResult(
        equity=pd.Series(eq, index=idx),
        returns=pd.Series(pd.Series(eq, index=idx).pct_change().fillna(0.0), index=idx),
        allocation=pd.DataFrame(alloc, index=idx, columns=p.columns),
        turnover=pd.Series(turns, index=idx),
        costs=pd.Series(costs, index=idx),
        rebalance_dates=sorted(rebal_set & set(idx)),
    )


def parse_weights(spec: str) -> dict[str, float]:
    """Parse ONE symbol or ONE portfolio spec. Returns {symbol: weight}.

    Commas separate the elements of the portfolio; each element is
    'SYM' (equal weight) or 'SYM:w' (weight w > 0). A bare 'SYM' on its
    own is a single symbol, not a portfolio. Surrounding whitespace is
    ignored. For specs that may hold several space-separated
    symbols/portfolios, use parse_items().
    """
    parts = [t.strip() for t in spec.strip().split(",") if t.strip()]
    if not parts:
        raise ValueError("no symbols given")
    out: dict[str, float] = {}
    for token in parts:
        if ":" in token:
            sym, w_s = token.rsplit(":", 1)
            try:
                w = float(w_s)
            except ValueError:
                raise ValueError(f"bad weight in '{token}'") from None
            if not math.isfinite(w) or w <= 0:
                raise ValueError(f"weight in '{token}' must be a positive number")
        else:
            sym, w = token, 1.0
        sym = sym.strip()
        if not sym:
            raise ValueError(f"empty symbol in '{token}'")
        out[sym] = out.get(sym, 0.0) + w
    if not out:
        raise ValueError("no symbols parsed")
    return out


def parse_items(spec: str) -> list[dict[str, float]]:
    """Parse a spec that may hold several symbols/portfolios.

    Spaces (or newlines) separate DISTINCT symbols/portfolios; commas
    join the elements of one portfolio. Each item is parsed by
    parse_weights().
    """
    return [parse_weights(item) for item in spec.split() if item.strip()]