import numpy as np
import pandas as pd
from scipy.optimize import minimize

def calculate_equilibrium_returns(weights, cov_matrix, risk_aversion=2.5):
    """
    逆向优化：从权重反推预期收益率
    weights: 市场组合权重（如沪深300成分股权重）
    cov_matrix: 收益率协方差矩阵
    risk_aversion: 风险厌恶系数（通常2.5-3.0）
    """
    # Pi = risk_aversion * Sigma * w
    pi = risk_aversion * np.dot(cov_matrix, weights)
    
    return pi

# 示例：假设我们有5只股票
stocks = ['股票A', '股票B', '股票C', '股票D', '股票E']
n_assets = len(stocks)

# 市场权重（模拟沪深300成分股权重）
market_weights = np.array([0.3, 0.25, 0.2, 0.15, 0.1])

# 协方差矩阵（模拟）
cov_matrix = np.array([
    [0.04, 0.02, 0.01, 0.005, 0.002],
    [0.02, 0.05, 0.015, 0.008, 0.003],
    [0.01, 0.015, 0.03, 0.01, 0.004],
    [0.005, 0.008, 0.01, 0.035, 0.006],
    [0.002, 0.003, 0.004, 0.006, 0.025]
])

# 计算均衡隐含收益率
risk_aversion = 2.5
pi = calculate_equilibrium_returns(market_weights, cov_matrix, risk_aversion)

print("均衡隐含收益率：")
for i, stock in enumerate(stocks):
    print(f"{stock}: {pi[i]:.4f} ({pi[i]*100:.2f}%)")

# 观点矩阵 P (n_views x n_assets)
# 每一行代表一个观点，每一列代表一只股票

# 示例：两个观点
# 观点1：股票A年化收益15%
# 观点2：股票A比股票B多涨5%

P = np.array([
    [1, 0, 0, 0, 0],   # 观点1：100%配置在股票A
    [1, -1, 0, 0, 0]   # 观点2：股票A - 股票B = 5%
])

# 观点收益向量 Q (n_views x 1)
Q = np.array([0.15, 0.05])  # 15% 和 5%

# 观点不确定性 Omega (n_views x n_views)
# 对角线元素表示每个观点的置信度（越小越自信）
tau = 0.05
omega = np.diag(np.diag(P @ (tau * cov_matrix) @ P.T))

print("观点矩阵 P：")
print(P)
print("\n观点收益 Q：")
print(Q)
print("\n观点不确定性 Omega：")
print(omega)

def black_litterman(pi, cov_matrix, P, Q, tau=0.05):
    """
    Black-Litterman模型主函数
    pi: 均衡隐含收益率
    cov_matrix: 协方差矩阵
    P: 观点矩阵
    Q: 观点收益向量
    tau: 缩放因子
    """
    n = len(pi)
    
    # 计算观点不确定性 Omega
    omega = np.diag(np.diag(P @ (tau * cov_matrix) @ P.T))
    omega_inv = np.linalg.inv(omega)
    
    # BL公式
    M_inv = np.linalg.inv(tau * cov_matrix) + P.T @ omega_inv @ P
    M = np.linalg.inv(M_inv)
    
    posterior_return = M @ (np.linalg.inv(tau * cov_matrix) @ pi + P.T @ omega_inv @ Q)
    
    # 后验协方差矩阵
    posterior_cov = M @ (np.linalg.inv(tau * cov_matrix) @ cov_matrix @ np.linalg.inv(tau * cov_matrix) + P.T @ omega_inv @ P) @ M
    
    return posterior_return, posterior_cov

# 计算后验收益率
posterior_return, posterior_cov = black_litterman(pi, cov_matrix, P, Q, tau)

print("\n后验预期收益率（BL调整后）：")
for i, stock in enumerate(stocks):
    print(f"{stock}: {posterior_return[i]:.4f} ({posterior_return[i]*100:.2f}%)")

print("\n后验协方差矩阵：")
print(posterior_cov)

def portfolio_optimization(expected_returns, cov_matrix, risk_aversion=2.5):
    """
    均值方差优化
    """
    n = len(expected_returns)
    
    # 目标函数：最大化效用 = 预期收益 - 0.5 * risk_aversion * 风险
    def objective(weights):
        portfolio_return = np.dot(weights, expected_returns)
        portfolio_risk = np.dot(weights, np.dot(cov_matrix, weights))
        utility = portfolio_return - 0.5 * risk_aversion * portfolio_risk
        return -utility  # scipy最小化，所以加负号
    
    # 约束：权重和为1
    constraints = ({'type': 'eq', 'fun': lambda x: np.sum(x) - 1})
    
    # 边界：不允许做空（可选）
    bounds = tuple((0, 1) for _ in range(n))
    
    # 初始权重：等权重
    initial_weights = np.array([1/n] * n)
    
    # 优化
    result = minimize(objective, initial_weights, method='SLSQP', 
                     bounds=bounds, constraints=constraints)
    
    return result.x

# 用BL后验收益率优化
optimal_weights_bl = portfolio_optimization(posterior_return, posterior_cov, risk_aversion)

print("\nBL模型最优权重：")
for i, stock in enumerate(stocks):
    print(f"{stock}: {optimal_weights_bl[i]:.4f} ({optimal_weights_bl[i]*100:.2f}%)")

# 对比：用历史收益率优化（传统Markowitz）
historical_returns = np.array([0.12, 0.10, 0.08, 0.06, 0.04])  # 模拟
optimal_weights_mv = portfolio_optimization(historical_returns, cov_matrix, risk_aversion)

print("\n传统Markowitz最优权重：")
for i, stock in enumerate(stocks):
    print(f"{stock}: {optimal_weights_mv[i]:.4f} ({optimal_weights_mv[i]*100:.2f}%)")

# 构建观点矩阵
stocks = top10_stocks.index.tolist()
n = len(stocks)

P = []
Q = []

# 观点1：新能源跑赢大盘10%
ev_stocks = ['300750.SZ', '002594.SZ']  # 宁德时代、比亚迪
ev_idx = [stocks.index(s) for s in ev_stocks]
p1 = np.zeros(n)
p1[ev_idx] = 1 / len(ev_idx)  # 等权重平均
P.append(p1)
Q.append(0.10)

# 观点2：金融板块持平（相对观点，略）
# 观点3：茅台15%
moutai_idx = stocks.index('600519.SH')
p3 = np.zeros(n)
p3[moutai_idx] = 1
P.append(p3)
Q.append(0.15)

P = np.array(P)
Q = np.array(Q)

# 市场权重
market_weights = top10_stocks.values

# 均衡隐含收益率
pi = calculate_equilibrium_returns(market_weights, cov_matrix, risk_aversion=2.5)

# BL后验收益率
posterior_return, posterior_cov = black_litterman(pi, cov_matrix, P, Q, tau=0.05)

# 优化
optimal_weights = portfolio_optimization(posterior_return, posterior_cov, risk_aversion=2.5)

# 对比市场权重
comparison = pd.DataFrame({
    '市场权重': market_weights,
    'BL最优权重': optimal_weights
}, index=stocks)

print("\n市场权重 vs BL最优权重：")
print(comparison)
print(f"\n权重变化（BL - 市场）：")
print(comparison['BL最优权重'] - comparison['市场权重'])

# 宏观因子观点
macro_factors = ['利率', '通胀', 'GDP']
F = np.array([
    [0.5, 0.3, 0.2],  # 利率对股票收益的影响
    [-0.3, 0.6, 0.1], # 通胀的影响
    [0.4, 0.2, 0.7]   # GDP的影响
])

# 宏观观点：我认为利率会下降1%
macro_view = np.array([-0.01, 0, 0])  # 利率-1%，通胀和GDP不变

# 将宏观观点转换为股票收益观点
Q_macro = F @ macro_view
P_macro = np.eye(n_assets)  # 对每个股票都有影响

# 结合微观和宏观观点
P_combined = np.vstack([P_micro, P_macro])
Q_combined = np.hstack([Q_micro, Q_macro])