import torch
import torch.nn.functional as F
import numpy as np
from typing import Optional, Callable

class MPPIPlanner:
    def __init__(self, model, n_assets: int, horizon=20, n_samples=200, n_iterations=20,
                 temperature=0.005, action_std=0.2, signal_prior=0.33):
        self.model = model
        self.n_assets = n_assets
        self.horizon = horizon
        self.n_samples = n_samples
        self.n_iterations = n_iterations
        self.temperature = temperature
        self.action_std = action_std
        self.signal_prior = signal_prior

    def plan(self, state_series, cost_fn=None, initial_weights=None):
        device = state_series.device
        B = 1
        if initial_weights is None:
            mean_weights = torch.ones(B, self.horizon, self.n_assets, device=device) / self.n_assets
        else:
            mean_weights = initial_weights.unsqueeze(0).expand(B, self.horizon, self.n_assets)
        std_weights = torch.ones_like(mean_weights) * self.action_std
        mean_signals = torch.zeros(B, self.horizon, self.n_assets, device=device, dtype=torch.long)
        for it in range(self.n_iterations):
            noise = torch.randn(B * self.n_samples, self.horizon, self.n_assets, device=device)
            sampled_weights = mean_weights.unsqueeze(1) + std_weights.unsqueeze(1) * noise
            sampled_weights = F.softmax(sampled_weights, dim=-1)
            sampled_signals = torch.randint(0, 3, (B * self.n_samples, self.horizon, self.n_assets), device=device)
            costs = torch.zeros(B * self.n_samples, device=device)
            current_state = state_series.repeat(self.n_samples, 1, 1)
            for t in range(self.horizon):
                w_t = sampled_weights[:, t]
                s_t = sampled_signals[:, t]
                z_next = self.model.predict_next_state(current_state, w_t, s_t)
                if cost_fn is not None:
                    c_t = cost_fn(z_next)
                else:
                    c_t = -z_next.mean(dim=(1, 2))
                costs += c_t
            costs = costs.reshape(B, self.n_samples)
            beta = costs.min(dim=1, keepdim=True).values
            weights_mppi = torch.exp(-(costs - beta) / self.temperature)
            weights_mppi = weights_mppi / weights_mppi.sum(dim=1, keepdim=True)
            flat_weights = sampled_weights.reshape(B, self.n_samples, self.horizon, self.n_assets)
            mean_weights = (weights_mppi[:, :, None, None] * flat_weights).sum(dim=1)
            std_weights = torch.sqrt((weights_mppi[:, :, None, None] * (flat_weights - mean_weights.unsqueeze(1)) ** 2).sum(dim=1) + 1e-4)
        best_idx = costs.argmin(dim=1)
        best_weights = sampled_weights.reshape(B, self.n_samples, self.horizon, self.n_assets)[torch.arange(B), best_idx, 0]
        best_signals = sampled_signals.reshape(B, self.n_samples, self.horizon, self.n_assets)[torch.arange(B), best_idx, 0]
        return {"weights": best_weights, "signals": best_signals, "expected_cost": costs.min(dim=1).values}

class CEMPlanner:
    def __init__(self, model, n_assets: int, horizon=20, n_candidates=200, n_elites=20, n_iterations=5, action_std=0.3):
        self.model = model
        self.n_assets = n_assets
        self.horizon = horizon
        self.n_candidates = n_candidates
        self.n_elites = n_elites
        self.n_iterations = n_iterations
        self.action_std = action_std

    def plan(self, state_series, cost_fn=None):
        device = state_series.device
        B = state_series.size(0)
        mean_weights = torch.ones(B, self.horizon, self.n_assets, device=device) / self.n_assets
        std_weights = torch.ones_like(mean_weights) * self.action_std
        for it in range(self.n_iterations):
            noise = torch.randn(B, self.n_candidates, self.horizon, self.n_assets, device=device)
            candidates = mean_weights.unsqueeze(1) + std_weights.unsqueeze(1) * noise
            candidates = F.softmax(candidates, dim=-1)
            signals = torch.randint(0, 3, (B, self.n_candidates, self.horizon, self.n_assets), device=device)
            costs = torch.zeros(B, self.n_candidates, device=device)
            for t in range(self.horizon):
                w_t = candidates[:, :, t].reshape(-1, self.n_assets)
                s_t = signals[:, :, t].reshape(-1, self.n_assets)
                state_rep = state_series.unsqueeze(1).repeat(1, self.n_candidates, 1, 1).reshape(-1, state_series.size(1), state_series.size(2))
                z_next = self.model.predict_next_state(state_rep, w_t, s_t)
                if cost_fn is not None:
                    c = cost_fn(z_next)
                else:
                    c = -z_next.mean(dim=(1, 2))
                costs += c.reshape(B, self.n_candidates)
            _, elite_idx = torch.topk(costs, self.n_elites, dim=1, largest=False)
            elite_weights = torch.gather(candidates, 1,
                elite_idx.unsqueeze(-1).unsqueeze(-1).expand(-1, -1, self.horizon, self.n_assets))
            mean_weights = elite_weights.mean(dim=1)
            std_weights = elite_weights.std(dim=1) + 1e-4
        best_idx = costs.argmin(dim=1)
        best_weights = candidates[torch.arange(B), best_idx, 0]
        best_signals = signals[torch.arange(B), best_idx, 0]
        return {"weights": best_weights, "signals": best_signals, "expected_cost": costs.min(dim=1).values}

def sharpe_cost(z_pred, target_return=0.0):
    returns = z_pred.mean(dim=(1, 2))
    return -returns