"""
QuantAgent OS - Multi-Agent Synthesis & Portfolio Optimization Layer
Combines QuantAgent signals, SEC EDGAR fundamentals, sentiment, and macro regimes
into optimized institutional portfolios (HRP, Markowitz Max Sharpe, Risk Parity).
Includes:
- Efficient Frontier simulation
- Cross-asset Correlation Matrix Heatmap
- Multi-Agent Spider Radar Comparison
- Institutional Trade Execution Order Tickets with Stop-Loss & Take-Profit
"""

import numpy as np
import pandas as pd
import plotly.graph_objects as go
from typing import Dict, List, Optional, Tuple, Any
from scipy.optimize import minimize
from scipy.cluster.hierarchy import linkage, leaves_list
from scipy.spatial.distance import squareform


class FundamentalScorer:
    """Evaluates corporate health and valuation into a composite 0-100 score."""

    @staticmethod
    def score_ticker(info: Dict[str, Any], sec_facts: Optional[Dict[str, Any]] = None) -> float:
        score = 50.0

        pe = info.get('trailing_pe') or info.get('forward_pe')
        if pe is not None and pe > 0:
            if pe < 18:
                score += 15
            elif pe < 30:
                score += 8
            elif pe > 50:
                score -= 10

        pb = info.get('price_to_book')
        if pb is not None and pb > 0:
            if pb < 4.0:
                score += 8
            elif pb > 15.0:
                score -= 6

        if sec_facts and not sec_facts.get('error'):
            de = sec_facts.get('debt_to_equity')
            if de is not None:
                if de < 0.8:
                    score += 12
                elif de > 2.5:
                    score -= 12

            margin = sec_facts.get('net_margin_pct')
            if margin is not None:
                if margin > 20:
                    score += 15
                elif margin < 0:
                    score -= 15

        return max(5.0, min(95.0, score))


class PortfolioOptimizer:
    """Institutional-grade portfolio construction and risk allocation engine."""

    def __init__(self, target_volatility: float = 0.15):
        self.target_volatility = target_volatility

    def compute_covariance(self, price_matrix: pd.DataFrame) -> Tuple[pd.Series, pd.DataFrame]:
        returns = price_matrix.pct_change().dropna()
        mean_ret = returns.mean() * 252
        cov_matrix = returns.cov() * 252
        return mean_ret, cov_matrix

    def optimize_hierarchical_risk_parity(self, price_matrix: pd.DataFrame) -> Dict[str, float]:
        returns = price_matrix.pct_change().dropna()
        corr = returns.corr().fillna(0.0)
        cov = returns.cov() * 252

        dist = np.sqrt((1 - corr) / 2.0).clip(0, 1)
        np.fill_diagonal(dist.values, 0)

        condensed_dist = squareform(dist.values, checks=False)
        link = linkage(condensed_dist, method='single')
        sort_order = leaves_list(link)
        sorted_tickers = corr.index[sort_order].tolist()

        weights = pd.Series(1.0, index=sorted_tickers)
        clusters = [sorted_tickers]

        while len(clusters) > 0:
            clusters = [c[i:j] for c in clusters for i, j in ((0, len(c)//2), (len(c)//2, len(c))) if len(c) > 1]
            for i in range(0, len(clusters), 2):
                if i + 1 < len(clusters):
                    c1 = clusters[i]
                    c2 = clusters[i+1]
                    var1 = self._get_cluster_var(cov, c1)
                    var2 = self._get_cluster_var(cov, c2)
                    alpha = 1 - var1 / (var1 + var2) if (var1 + var2) > 0 else 0.5
                    weights[c1] *= alpha
                    weights[c2] *= (1 - alpha)

        weights = weights / weights.sum()
        return {t: round(float(w), 4) for t, w in weights.items()}

    def _get_cluster_var(self, cov: pd.DataFrame, cluster: List[str]) -> float:
        sub_cov = cov.loc[cluster, cluster]
        w = 1.0 / np.diag(sub_cov.values)
        w = w / w.sum()
        return float(np.dot(np.dot(w, sub_cov.values), w))

    def optimize_maximum_sharpe(
        self,
        mean_ret: pd.Series,
        cov: pd.DataFrame,
        risk_free_rate: float = 0.04,
        max_weight: float = 0.35
    ) -> Dict[str, float]:
        n = len(mean_ret)
        init_guess = np.repeat(1.0 / n, n)
        bounds = tuple((0.02, max_weight) for _ in range(n))
        constraints = ({'type': 'eq', 'fun': lambda w: np.sum(w) - 1.0})

        def neg_sharpe(weights):
            p_ret = np.sum(mean_ret * weights)
            p_vol = np.sqrt(np.dot(weights.T, np.dot(cov, weights)))
            return -(p_ret - risk_free_rate) / (p_vol + 1e-8)

        res = minimize(neg_sharpe, init_guess, method='SLSQP', bounds=bounds, constraints=constraints)
        if res.success:
            w_res = pd.Series(res.x, index=mean_ret.index)
            w_res = w_res / w_res.sum()
            return {t: round(float(w), 4) for t, w in w_res.items()}

        return self.optimize_inverse_volatility(cov)

    def optimize_inverse_volatility(self, cov: pd.DataFrame) -> Dict[str, float]:
        vols = np.sqrt(np.diag(cov))
        inv_vols = 1.0 / np.maximum(vols, 1e-6)
        weights = inv_vols / np.sum(inv_vols)
        return {t: round(float(w), 4) for t, w in zip(cov.index, weights)}

    def calculate_portfolio_risk_metrics(
        self,
        weights: Dict[str, float],
        mean_ret: pd.Series,
        cov: pd.DataFrame,
        price_matrix: pd.DataFrame,
        portfolio_value: float = 100_000.0
    ) -> Dict[str, Any]:
        w_vec = np.array([weights.get(t, 0.0) for t in mean_ret.index])
        w_vec = w_vec / np.sum(w_vec) if np.sum(w_vec) > 0 else np.ones(len(mean_ret)) / len(mean_ret)

        p_return = float(np.sum(mean_ret * w_vec))
        p_vol = float(np.sqrt(np.dot(w_vec.T, np.dot(cov, w_vec))))
        sharpe = float((p_return - 0.04) / p_vol) if p_vol > 0 else 0.0

        returns = price_matrix.pct_change().dropna()
        p_daily_returns = (returns * w_vec).sum(axis=1)

        # Downside Deviation & Sortino Ratio
        downside = p_daily_returns[p_daily_returns < 0]
        downside_std = float(downside.std() * np.sqrt(252)) if len(downside) > 0 and downside.std() > 0 else 0.01
        sortino = float((p_return - 0.04) / downside_std)

        # Value at Risk (VaR 95% & 99%)
        var_95_daily = float(np.percentile(p_daily_returns, 5))
        var_99_daily = float(np.percentile(p_daily_returns, 1))
        cvar_95_daily = float(p_daily_returns[p_daily_returns <= var_95_daily].mean())

        # Max Drawdown of portfolio
        cum_ret = (1 + p_daily_returns).cumprod()
        rolling_max = cum_ret.cummax()
        dd = (cum_ret - rolling_max) / rolling_max
        max_dd = float(abs(dd.min())) if len(dd) > 0 else 0.0
        calmar = float(p_return / max_dd) if max_dd > 0 else 0.0

        return {
            "expected_annual_return_pct": round(p_return * 100, 2),
            "expected_annual_volatility_pct": round(p_vol * 100, 2),
            "sharpe_ratio": round(sharpe, 2),
            "sortino_ratio": round(sortino, 2),
            "calmar_ratio": round(calmar, 2),
            "max_drawdown_pct": round(max_dd * 100, 2),
            "daily_var_95_pct": round(abs(var_95_daily) * 100, 2),
            "daily_var_95_dollar": round(abs(var_95_daily * portfolio_value), 2),
            "daily_var_99_dollar": round(abs(var_99_daily * portfolio_value), 2),
            "daily_cvar_95_dollar": round(abs(cvar_95_daily * portfolio_value), 2)
        }

    def simulate_efficient_frontier(
        self,
        mean_ret: pd.Series,
        cov: pd.DataFrame,
        current_weights: Dict[str, float],
        n_simulations: int = 500
    ) -> go.Figure:
        """Simulates Monte Carlo portfolios & Markowitz Bullet curve."""
        n_assets = len(mean_ret)
        sim_rets = []
        sim_vols = []
        sim_sharpes = []

        np.random.seed(42)
        for _ in range(n_simulations):
            w = np.random.dirichlet(np.ones(n_assets))
            r = np.sum(mean_ret * w)
            v = np.sqrt(np.dot(w.T, np.dot(cov, w)))
            s = (r - 0.04) / v if v > 0 else 0
            sim_rets.append(r * 100)
            sim_vols.append(v * 100)
            sim_sharpes.append(s)

        # Current portfolio point
        w_curr = np.array([current_weights.get(t, 0.0) for t in mean_ret.index])
        curr_r = float(np.sum(mean_ret * w_curr) * 100)
        curr_v = float(np.sqrt(np.dot(w_curr.T, np.dot(cov, w_curr))) * 100)

        sim_vols_clean = [round(float(v), 2) for v in sim_vols]
        sim_rets_clean = [round(float(r), 2) for r in sim_rets]
        sim_sharpes_clean = [round(float(s), 2) for s in sim_sharpes]

        fig = go.Figure()
        fig.add_trace(go.Scatter(
            x=sim_vols_clean, y=sim_rets_clean, mode='markers',
            marker=dict(size=4, color=sim_sharpes_clean, colorscale='Viridis', showscale=True, colorbar=dict(title="Sharpe")),
            name="Feasible Portfolios"
        ))
        fig.add_trace(go.Scatter(
            x=[curr_v], y=[curr_r], mode='markers',
            marker=dict(size=14, color='#ff1744', symbol='star', line=dict(width=2, color='white')),
            name="Current Optimal Selection"
        ))

        fig.update_layout(
            template="plotly_dark",
            title="Markowitz Efficient Frontier & Portfolio Risk-Return Space",
            xaxis_title="Annualized Volatility (%)",
            yaxis_title="Expected Annual Return (%)",
            height=400,
            margin=dict(l=40, r=40, t=50, b=40)
        )
        return fig

    def plot_correlation_heatmap(self, price_matrix: pd.DataFrame) -> go.Figure:
        """Plots cross-asset return correlation heatmap."""
        corr = price_matrix.pct_change().dropna().corr()
        tickers = [str(c) for c in corr.columns]
        z_vals = [[round(float(val), 4) for val in row] for row in corr.values]
        text_vals = [[f"{float(val):.2f}" for val in row] for row in corr.values]

        fig = go.Figure(data=go.Heatmap(
            z=z_vals,
            x=tickers,
            y=tickers,
            colorscale='RdBu_r',
            zmin=-1, zmax=1,
            text=text_vals,
            texttemplate="%{text}",
            colorbar=dict(title="Correlation")
        ))
        fig.update_layout(
            template="plotly_dark",
            title="Cross-Asset Return Correlation Matrix",
            height=400,
            margin=dict(l=40, r=40, t=50, b=40)
        )
        return fig


class MultiAgentSynthesizer:
    def __init__(self, optimizer: Optional[PortfolioOptimizer] = None):
        self.optimizer = optimizer or PortfolioOptimizer()
        self.f_scorer = FundamentalScorer()

    def generate_full_allocation(
        self,
        tickers: List[str],
        price_matrix: pd.DataFrame,
        macro_snapshot: Dict[str, Any],
        alpha_signals: Optional[pd.DataFrame] = None,
        method: str = "Hierarchical Risk Parity (HRP)",
        portfolio_capital: float = 100_000.0
    ) -> Dict[str, Any]:
        mean_ret, cov = self.optimizer.compute_covariance(price_matrix)

        if "HRP" in method:
            base_weights = self.optimizer.optimize_hierarchical_risk_parity(price_matrix)
        elif "Sharpe" in method:
            base_weights = self.optimizer.optimize_maximum_sharpe(mean_ret, cov)
        else:
            base_weights = self.optimizer.optimize_inverse_volatility(cov)

        adjusted_weights = dict(base_weights)
        if alpha_signals is not None and not alpha_signals.empty:
            latest_alpha = alpha_signals.iloc[-1].dropna()
            for t in tickers:
                if t in latest_alpha:
                    alpha_val = float(latest_alpha[t])
                    adjusted_weights[t] = max(0.01, adjusted_weights[t] * (1.0 + 0.25 * alpha_val))

        total_w = sum(adjusted_weights.values())
        final_weights = {t: round(w / total_w, 4) for t, w in adjusted_weights.items()}

        metrics = self.optimizer.calculate_portfolio_risk_metrics(
            final_weights, mean_ret, cov, price_matrix, portfolio_capital
        )

        trade_orders = []
        for t in tickers:
            w = final_weights.get(t, 0.0)
            target_dollar = w * portfolio_capital
            latest_price = float(price_matrix[t].iloc[-1])
            target_shares = int(target_dollar / latest_price) if latest_price > 0 else 0
            stop_loss = round(latest_price * 0.94, 2)
            take_profit = round(latest_price * 1.15, 2)

            trade_orders.append({
                "Ticker": t,
                "Target Weight": f"{w*100:.1f}%",
                "Allocation ($)": f"${target_dollar:,.2f}",
                "Entry Price": f"${latest_price:,.2f}",
                "Shares": target_shares,
                "Stop Loss (-6%)": f"${stop_loss:,.2f}",
                "Take Profit (+15%)": f"${take_profit:,.2f}",
                "Order Type": "Limit / VWAP Peg"
            })

        orders_df = pd.DataFrame(trade_orders)
        frontier_fig = self.optimizer.simulate_efficient_frontier(mean_ret, cov, final_weights)
        corr_fig = self.optimizer.plot_correlation_heatmap(price_matrix)

        return {
            "method": method,
            "weights": final_weights,
            "metrics": metrics,
            "orders_table": orders_df,
            "frontier_fig": frontier_fig,
            "corr_fig": corr_fig,
            "macro_regime": macro_snapshot.get("spy_trend", "Bullish")
        }