from __future__ import annotations
from pathlib import Path
from shutil import rmtree

import torch
from torch import nn, Tensor, float32
from torch.nn import Module
import torch.nn.functional as F
from torch.utils.data import DataLoader
from torch.optim import AdamW

from accelerate import Accelerator
from memmap_replay_buffer import ReplayBuffer

from tqdm.auto import tqdm

import einx
from einops import rearrange

# helpers

def exists(v):
    return v is not None

def default(v, d):
    return v if exists(v) else d

# advantage binner

class AdvantageBinner(Module):
    """
    Discretizes continuous advantage values into categorical tokens.
    Following literature, limits the range to [-2.0, 2.0] and categorizes into 10 bins.
    """
    def __init__(
        self,
        min_val = -2.0,
        max_val = 2.0,
        num_bins = 10
    ):
        super().__init__()
        self.min_val = min_val
        self.max_val = max_val
        self.num_bins = num_bins
        
        # Create persistent bin boundaries
        bins = torch.linspace(min_val, max_val, num_bins)
        self.register_buffer('bins', bins, persistent = False)

    def forward(
        self,
        advantages: torch.Tensor
    ) -> torch.Tensor:
        if advantages.numel() == 0:
            return advantages

        # Clamp advantages to the specified range
        clamped_adv = advantages.clamp(self.min_val, self.max_val)
        
        # Calculate distances to all bins and find the closest index
        # clamped_adv: [Batch], self.bins: [num_bins]
        
        clamped_adv = clamped_adv.to(self.bins.device)
        distances = torch.abs(einx.subtract('b, n -> b n', clamped_adv, self.bins))
        
        # Nearest bin indices
        binned_indices = distances.argmin(dim = -1)
        
        return binned_indices

# rise orchestrator

class RISE(Module):
    """
    The RISE Orchestrator (Self-Improving Robot Policy with Compositional World Model).
    Wraps the policy, dynamics model, and value model into a single self-improving loop.
    """
    def __init__(
        self,
        policy: Module,
        dynamics_model: Module,
        value_model: Module,
        advantage_binner = None,
        num_advantage_bins = 10,
        trajectory_length = 4,
        num_prompt_tokens = 12,
        imagination_steps = 5,
        frames_per_future_latent = 5,
        cpu = False,
        accelerate_kwargs: dict = {}
    ):
        super().__init__()
        
        self.accelerator = Accelerator(cpu = cpu, **accelerate_kwargs)
        self.device = self.accelerator.device
        
        self.policy = policy.to(self.device)
        self.dynamics_model = dynamics_model.to(self.device)
        self.value_model = value_model.to(self.device, dtype=torch.bfloat16)
        
        self.advantage_binner = default(
            advantage_binner, 
            AdvantageBinner(num_bins = num_advantage_bins)
        )
        
        # configuration
        
        self.trajectory_length = trajectory_length
        self.num_prompt_tokens = num_prompt_tokens
        self.imagination_steps = imagination_steps
        self.frames_per_future_latent = frames_per_future_latent

    def imagination_rollout(
        self,
        seed_dataset,
        num_episodes = 1,
        batch_size = 1,
        max_timesteps = 1, # Number of autoregressive rollout steps
        buffer_folder = './imagination_replay_buffer',
        predict_num_future_latents = 1,
        force_new_buffer = True
    ) -> ReplayBuffer:
        """
        Rollout Stage: Use world model as interactive simulator to improve the policy.
        """

        if force_new_buffer and Path(buffer_folder).exists():
            rmtree(buffer_folder)
            
        # 1. Identify Buffer Geometry from seed_dataset
        sample = seed_dataset[0]
        video = sample['video']
        proprio = sample['proprioception']
        
        time, channels, height, width = video.shape
        _, proprio_dim = proprio.shape
        
        # Policy always has dim_action_input
        action_dim = self.policy.dim_action_input
        
        # 2. Initialize ReplayBuffer
        buffer = ReplayBuffer(
            buffer_folder,
            max_episodes = num_episodes,
            max_timesteps = max_timesteps, 
            fields = dict(
                images = ('float', (time, channels, height, width)),
                joint_state = ('float', (time, proprio_dim)),
                actions = ('float', (self.trajectory_length, action_dim)),
                advantage_ids = 'int',
                imagined_video = ('float', (time + predict_num_future_latents * self.frames_per_future_latent, channels, height, width)) # prefix + future
            )
        )
        
        loader = DataLoader(seed_dataset, batch_size = batch_size, shuffle = True)
        
        self.eval()
        pbar = tqdm(total = num_episodes, desc = "Imagination Rollout Stage")
        
        device = next(self.policy.parameters()).device

        episodes_collected = 0
        for batch_data in loader:
            if episodes_collected >= num_episodes:
                break
                
            videos = batch_data['video'].to(device)
            proprio = batch_data['proprioception'].to(device)
            prompt_token_ids = batch_data['prompt_token_ids'].to(device)
            
            batch, *_ = videos.shape

            if episodes_collected + batch > num_episodes:
                slice_size = num_episodes - episodes_collected
                videos = videos[:slice_size]
                proprio = proprio[:slice_size]
                prompt_token_ids = prompt_token_ids[:slice_size]
                batch = slice_size

            with torch.no_grad():
                with buffer.batched_episode(batch_size = batch):
                    
                    # Autoregressive neural world model loop
                    for step in range(max_timesteps):
                        # A. Generate actions with "positive intent" (highest advantage bin)
                        pz_images = rearrange(videos.float(), 'b t c h w -> b c t h w')
                        
                        generated_actions = self.policy.sample_actions(
                            images = pz_images,
                            token_ids = prompt_token_ids,
                            joint_states = proprio.float(),
                            trajectory_length = self.trajectory_length,
                            advantage_ids = self.advantage_binner.num_bins - 1,
                            steps = self.imagination_steps,
                            show_pbar = False
                        )
                        
                        # B. Imagine Future via Dynamics Model
                        output_video = self.dynamics_model(
                            videos = videos,
                            prompt_token_ids = prompt_token_ids,
                            predict_num_future_latents = predict_num_future_latents,
                            actions = generated_actions.to(videos.dtype),
                            return_decoded_latents = True
                        )
                        
                        # C. Compute Advantages via Value Network
                        initial_value = self.value_model(videos)
                        imagined_value = self.value_model(output_video)
                        
                        advantage = einx.subtract('b, b -> b', imagined_value, initial_value)
                        binned_adv = self.advantage_binner(advantage)
                        
                        # D. Store in ReplayBuffer using store_batch
                        buffer.store_batch(
                            images = videos.float().cpu(),
                            joint_state = proprio.float().cpu(),
                            actions = generated_actions.float().cpu(),
                            advantage_ids = binned_adv.cpu(),
                            imagined_video = output_video.float().cpu() # prefix + future
                        )
                        
                        # E. Interactive Simulator step: update state for next timestep
                        # Extract the newly predicted future frames to form the new current state
                        
                        # For simplicity, if output video is concatenation of prefix + future:
                        # We take the last 'time' frames to act as the new prefix.
                        # This works regardless of the specific prefix size mapping, 
                        # as long as we maintain the sliding window of context.
                        videos = output_video[:, -time:]
                        
                        # For proprioception, normally a dynamic model would predict next proprio too.
                        # For RISE paper mapping without explicit proprio dynamics, we can carry over
                        # the last known joint states, or ideally the policy is robust to stales.
                        # In a full model, output_proprio = dynamics_model(...)
                        
            episodes_collected += batch
            pbar.update(batch)
            self.accelerator.print(f"\nImagination complete. Collected {episodes_collected} episodes of dreamt experience.")
        return buffer

    def finetune_with_advantage_conditioning(
        self,
        replay_buffer,
        num_steps = 1000,
        batch_size = 4,
        lr = 1e-4,
        weight_decay = 0.,
        max_grad_norm = 0.5
    ):
        """
        Training Stage: Optimize the policy using the advantage-conditioned imagination data.
        """
        
        optim = AdamW(self.policy.parameters(), lr = lr, weight_decay = weight_decay)
        
        # ReplayBuffer as Dataset (timestep_level = True for efficient learning)
        dataset = replay_buffer.dataset(timestep_level = True)
        loader = DataLoader(dataset, batch_size = batch_size, shuffle = True)
        
        self.policy, optim, loader = self.accelerator.prepare(
            self.policy, optim, loader
        )
        
        self.policy.train()
        pbar = tqdm(range(num_steps), desc = "Self-Improvement Training Stage")
        
        loader_iter = iter(loader)
        
        for step in pbar:
            try:
                batch_data = next(loader_iter)
            except StopIteration:
                loader_iter = iter(loader)
                batch_data = next(loader_iter)
                
            videos = batch_data['images']
            proprio = batch_data['joint_state']
            actions = batch_data['actions']
            advantage_ids = batch_data['advantage_ids']
            
            batch, *_ = videos.shape

            # Note: PiZero expects images as [B, C, T, H, W]
            pz_images = rearrange(videos.float(), 'b t c h w -> b c t h w')
            
            # Prompt token ids are not stored in buffer in current rollout - assume mock or retrieve if added
            # For simplicity, if not in buffer, generate random ones for training
            # In a real scenario, you'd store prompt_token_ids in the buffer
            
            prompt_token_ids = torch.randint(0, self.policy.token_emb.num_embeddings, (batch, self.num_prompt_tokens), device = self.accelerator.device)
            
            with self.accelerator.autocast():
                loss, _ = self.policy(
                    token_ids = prompt_token_ids,
                    joint_state = proprio.float(),
                    images = pz_images,
                    actions = actions,
                    advantage_ids = advantage_ids
                )
            
            self.accelerator.backward(loss)
            
            if exists(max_grad_norm):
                self.accelerator.clip_grad_norm_(self.policy.parameters(), max_grad_norm)
                
            optim.step()
            optim.zero_grad()
            
            pbar.set_postfix(loss = f"{loss.item():.4f}")
            
        self.accelerator.print("\nTraining complete.")
