import torch
import torch.nn as nn
import torch.optim as optim
from pylon.core import compile

# 定义LSTM模型
class LSTMModel(nn.Module):
    def __init__(self, input_size, hidden_size, num_layers, output_size):
        super(LSTMModel, self).__init__()
        self.lstm = nn.LSTM(input_size, hidden_size, num_layers, batch_first=True)
        self.fc = nn.Linear(hidden_size, output_size)

    def forward(self, x):
        self.lstm.flatten_parameters()
        outputs, (hidden, cell) = self.lstm(x)
        predictions = self.fc(hidden[-1])
        return predictions

# 超参数
input_size = 10  # 特征数量
hidden_size = 50  # LSTM隐藏层大小
num_layers = 2  # LSTM层数
output_size = 3  # 输出大小，对应3种资产的权重
sequence_length = 5  # 输入序列长度
batch_size = 1  # 批处理大小

# 实例化LSTM模型
lstm_model = LSTMModel(input_size, hidden_size, num_layers, output_size)

# 随机生成一些示例数据
x = torch.rand(batch_size, sequence_length, input_size)

# 定义组合的预期收益率和协方差矩阵
expected_returns = torch.tensor([0.05, 0.06, 0.07])
covariance_matrix = torch.tensor([[0.01, 0.001, 0.0],
                                   [0.001, 0.02, 0.002],
                                   [0, 0.002, 0.03]])
risk_free_rate = 0.02

# 定义目标函数，计算夏普比率
def objective(weights, returns, cov_matrix, risk_free):
    portfolio_return = torch.dot(weights, returns)
    portfolio_volatility = torch.sqrt(torch.dot(weights.T, torch.mm(cov_matrix, weights)))
    sharpe_ratio = (portfolio_return - risk_free) / portfolio_volatility
    return -sharpe_ratio # 我们取负值因为我们要最大化夏普比率

# 定义组合权重约束
def portfolio_constraint(weights):
    return torch.all(weights >= 0), torch.all(weights <= 1), torch.all(weights.sum() - 1 < 1e-5)

# 编译约束函数
constraint_loss = compile(portfolio_constraint, backend='pytorch')

# 定义优化器
optimizer = optim.Adam(lstm_model.parameters(), lr=0.01)

# 训练循环
for epoch in range(100): # 假设我们训练100个epoch
    # 前向传播
    optimizer.zero_grad()
    weights_pred = lstm_model(x)
    objective_loss = objective(weights_pred, expected_returns, covariance_matrix, risk_free_rate)
    constraint_penalty = constraint_loss(weights_pred)
    total_loss = objective_loss + constraint_penalty
    total_loss.backward()
    optimizer.step()

    if epoch % 10 == 0:
        print(f"Epoch {epoch}, Sharpe Ratio: {-objective_loss.item()}, Constraint Penalty: {constraint_penalty.item()}")

# 打印最终的权重和夏普比率
print(f"Predicted Portfolio Weights: {weights_pred.detach()}")
print(f"Sharpe Ratio: {-objective_loss.detach().item()}")