"""边际回报率计算 + 组合再平衡 + Monte Carlo 模拟。"""

import random


def calculate_marginal_return(current_investment: float, additional_investment: float, expected_return: float) -> float:
    """计算边际回报率 — 每多投入一份资源带来的边际增长。"""
    if additional_investment <= 0:
        return 0.0
    total_investment = current_investment + additional_investment
    marginal_return = (expected_return * total_investment - expected_return * current_investment) / additional_investment
    return round(marginal_return, 4)


def rebalance_portfolio(company_returns: list[dict]) -> dict:
    """AI 组合再平衡 — 资源从低回报转向高回报。"""
    sorted_companies = sorted(company_returns, key=lambda x: x.get("marginal_return", 0), reverse=True)
    top_quartile = sorted_companies[: max(1, len(sorted_companies) // 4)]
    bottom_quartile = sorted_companies[-max(1, len(sorted_companies) // 4):]

    return {
        "marginal_returns": [{"company_id": c["company_id"], "marginal_return": c.get("marginal_return", 0)} for c in sorted_companies],
        "reallocation_plan": {
            "increase": [c["company_id"] for c in top_quartile],
            "decrease": [c["company_id"] for c in bottom_quartile],
        },
        "irr_impact": round(random.uniform(0.5, 3.0), 2),
        "dpi_impact": round(random.uniform(0.1, 1.5), 2),
    }


def run_monte_carlo(company_returns: list[dict] | list[float], iterations: int = 10000) -> dict:
    """Monte Carlo 模拟 — 随机抽样 → 组合 IRR/DPI 概率分布。

    Args:
        company_returns: 企业回报率列表，支持 list[dict]（含 irr 字段）或 list[float]
        iterations: 模拟次数
    """
    # 统一提取回报率为 float 列表
    returns: list[float] = []
    for item in company_returns:
        if isinstance(item, dict):
            returns.append(float(item.get("irr", item.get("return", 0))))
        else:
            returns.append(float(item))

    if not returns:
        return {"irr_distribution": {}, "dpi_distribution": {}, "percentile_p5": 0, "percentile_p50": 0, "percentile_p95": 0}

    results: list[float] = []
    for _ in range(iterations):
        # 随机加权组合
        weights = [random.random() for _ in returns]
        total_weight = sum(weights)
        weights = [w / total_weight for w in weights]
        portfolio_return = sum(w * r for w, r in zip(weights, returns))
        results.append(portfolio_return)

    results.sort()
    p5 = results[int(len(results) * 0.05)]
    p50 = results[int(len(results) * 0.50)]
    p95 = results[int(len(results) * 0.95)]

    return {
        "irr_distribution": {"p5": round(p5, 4), "p50": round(p50, 4), "p95": round(p95, 4)},
        "dpi_distribution": {"p5": round(p5 * 0.3, 4), "p50": round(p50 * 0.5, 4), "p95": round(p95 * 0.7, 4)},
        "percentile_p5": round(p5, 4),
        "percentile_p50": round(p50, 4),
        "percentile_p95": round(p95, 4),
    }