"""
CFD Flow Control Environment using CelerisLab.

A Gymnasium environment for active flow control using lattice Boltzmann simulation.
"""

import os
from collections import deque
from typing import Optional, Tuple, Dict, Any

import gymnasium as gym
import numpy as np
from gymnasium import spaces

# Set threading to avoid conflicts with GPU
os.environ["OMP_NUM_THREADS"] = "1"
os.environ["MKL_NUM_THREADS"] = "1"


class CFDFlowControlEnv(gym.Env):
    """
    CFD flow control environment with cylinder and sensors.
    
    The environment simulates flow around a cylinder with multiple control cylinders
    and sensors to measure flow properties. The agent controls the cylinder velocities
    to optimize flow characteristics.
    
    Args:
        device_id: CUDA device ID to use for simulation
        config_cuda: CelerisLab CUDA configuration (optional, will load from config if None)
        config_field: CelerisLab flow field configuration (optional, will load from config if None)
        n_control_cylinders: Number of controllable cylinders (default: 3)
        n_sensors: Number of flow sensors (default: 3)
        max_steps: Maximum steps per episode (default: 500)
        sample_interval: Simulation steps between observations (default: 800)
        fifo_length: Length of state history (default: 120)
        convergence_length: Steps to check for convergence (default: 60)
        warmup_steps_factor: Multiple of grid size for warmup (default: 4)
    """
    
    metadata = {"render_modes": ["human"], "render_fps": 30}
    
    def __init__(
        self,
        device_id: int = 0,
        config_cuda = None,
        config_field = None,
        n_control_cylinders: int = 3,
        n_sensors: int = 3,
        max_steps: int = 500,
        sample_interval: int = 800,
        fifo_length: int = 120,
        convergence_length: int = 60,
        warmup_steps_factor: int = 4,
    ):
        super().__init__()
        
        # Load configurations if not provided
        if config_cuda is None or config_field is None:
            from ..config import load_celeris_configs
            config_cuda, config_field = load_celeris_configs()
        
        self.config_cuda = config_cuda
        self.config_field = config_field
        self.device_id = device_id
        
        # Environment parameters
        self.n_control = n_control_cylinders
        self.n_sensors = n_sensors
        self.max_steps = max_steps
        self.sample_interval = sample_interval
        self.fifo_length = fifo_length
        self.convergence_length = convergence_length
        self.warmup_steps_factor = warmup_steps_factor
        
        # Determine data type
        if config_field.data_type == "FP32":
            self.dtype = np.float32
        elif config_field.data_type == "FP64":
            self.dtype = np.float64
        else:
            raise ValueError(f"Unsupported data type: {config_field.data_type}")
        
        # Action and observation dimensions
        # Action: velocity control for n cylinders (x, y, rotation)
        self.action_dim = n_control_cylinders
        # Observation: sensor readings (u, v) from n sensors  
        self.obs_dim = n_sensors * 2 * 2  # 2 velocity components × 2 (current + derivative)
        
        # Gym spaces
        self.action_space = spaces.Box(
            low=-1.0,
            high=1.0,
            shape=(self.action_dim,),
            dtype=self.dtype
        )
        self.observation_space = spaces.Box(
            low=-np.inf,
            high=np.inf,
            shape=(self.obs_dim,),
            dtype=self.dtype
        )
        
        # State tracking
        self.fifo_states = deque(maxlen=fifo_length)
        self.target_states = np.empty((0, self.n_sensors * 2), dtype=self.dtype)
        self.current_step = 0
        
        # Normalization factors (will be set during warmup)
        self.sens_norm_fact = np.ones(self.n_sensors * 2, dtype=self.dtype)
        self.sens_deviation = np.zeros(self.n_sensors * 2, dtype=self.dtype)
        
        # Reward tracking
        self.reward_cd = 0.0
        self.reward_cl = 0.0
        self.reward_sim = 0.0
        
        # Initialize flow field
        self._init_flow_field()
    
    def _init_flow_field(self):
        """Initialize the CelerisLab flow field simulation."""
        from CelerisLab import FlowField
        
        self.flow_field = FlowField(
            self.config_field,
            self.config_cuda,
            self.device_id
        )
        
        # Get grid parameters
        L0 = 20  # Characteristic length
        U0 = self.config_field.velocity
        NX = self.flow_field.FIELD_SHAPE[0]
        NY = self.flow_field.FIELD_SHAPE[1]
        
        # Add main cylinder (obstacle)
        center = (10 * L0, (NY - 1) / 2, 0)
        self.flow_field.add_cylinder(center, L0)
        
        # Add sensors
        sensor_y_positions = [
            (NY - 1) / 2 + 2 * L0,  # Above centerline
            (NY - 1) / 2,            # At centerline  
            (NY - 1) / 2 - 2 * L0,  # Below centerline
        ]
        for i in range(min(self.n_sensors, len(sensor_y_positions))):
            center = (40 * L0, sensor_y_positions[i], 0)
            self.flow_field.add_sensor(center, L0 / 4)
        
        # Warmup simulation
        warmup_steps = int(self.warmup_steps_factor * NX / U0)
        self.flow_field.run(warmup_steps, np.zeros(self.n_control + 1, dtype=self.dtype))
        
        # Collect baseline states for normalization
        for _ in range(self.fifo_length):
            self.flow_field.run(
                self.sample_interval,
                np.zeros(self.n_control + 1, dtype=self.dtype)
            )
            new_state = self.flow_field.obs.copy()[2:2 + self.n_sensors * 2]
            self.target_states = np.vstack((self.target_states, new_state))
            self.fifo_states.append(new_state)
        
        # Calculate normalization factors
        self._calculate_normalization()
    
    def _calculate_normalization(self):
        """Calculate normalization factors from baseline states."""
        if len(self.target_states) > 0:
            self.sens_norm_fact = np.std(self.target_states, axis=0) + 1e-6
            self.sens_deviation = np.mean(self.target_states, axis=0)
    
    def _normalize_state(self, state: np.ndarray) -> np.ndarray:
        """Normalize state using calculated factors."""
        return (state - self.sens_deviation) / self.sens_norm_fact
    
    def _compute_reward(self, state: np.ndarray, action: np.ndarray) -> float:
        """
        Compute reward based on drag reduction and flow similarity.
        
        Args:
            state: Current state observation
            action: Applied action
        
        Returns:
            Total reward
        """
        # Get force measurements from simulation
        obs = self.flow_field.obs
        cd = obs[0]  # Drag coefficient
        cl = obs[1]  # Lift coefficient
        
        # Drag reduction reward (negative drag is good)
        self.reward_cd = -cd * 0.1
        
        # Lift minimization (want symmetric flow)
        self.reward_cl = -abs(cl) * 0.05
        
        # Flow similarity to baseline (want smooth control)
        if len(self.fifo_states) >= self.convergence_length:
            recent_states = np.array(list(self.fifo_states)[-self.convergence_length:])
            target_recent = self.target_states[-self.convergence_length:]
            
            # Dynamic Time Warping distance (simplified)
            diff = np.mean(np.abs(recent_states - target_recent))
            self.reward_sim = -diff * 0.5
        else:
            self.reward_sim = 0.0
        
        # Total reward
        total_reward = self.reward_cd + self.reward_cl + self.reward_sim
        
        return float(total_reward)
    
    def reset(
        self,
        seed: Optional[int] = None,
        options: Optional[Dict[str, Any]] = None
    ) -> Tuple[np.ndarray, Dict[str, Any]]:
        """
        Reset the environment to initial state.
        
        Args:
            seed: Random seed for reproducibility
            options: Additional options
        
        Returns:
            Tuple of (observation, info)
        """
        super().reset(seed=seed)
        
        self.current_step = 0
        self.fifo_states.clear()
        
        # Run a few steps to get initial state
        for _ in range(10):
            self.flow_field.run(
                self.sample_interval,
                np.zeros(self.n_control + 1, dtype=self.dtype)
            )
            state = self.flow_field.obs.copy()[2:2 + self.n_sensors * 2]
            self.fifo_states.append(state)
        
        # Get current state
        current_state = self.fifo_states[-1]
        
        # Compute state derivative (approximation)
        if len(self.fifo_states) >= 2:
            state_derivative = self.fifo_states[-1] - self.fifo_states[-2]
        else:
            state_derivative = np.zeros_like(current_state)
        
        # Normalize and concatenate
        obs = np.concatenate([
            self._normalize_state(current_state),
            self._normalize_state(state_derivative)
        ]).astype(self.dtype)
        
        info = {
            'step': self.current_step,
            'cd': self.flow_field.obs[0],
            'cl': self.flow_field.obs[1],
        }
        
        return obs, info
    
    def step(self, action: np.ndarray) -> Tuple[np.ndarray, float, bool, bool, Dict[str, Any]]:
        """
        Take a step in the environment.
        
        Args:
            action: Action to take (cylinder velocities)
        
        Returns:
            Tuple of (observation, reward, terminated, truncated, info)
        """
        # Convert action to control input
        # Action is in [-1, 1], scale to appropriate velocity range
        control = np.zeros(self.n_control + 1, dtype=self.dtype)
        control[:self.n_control] = action * 0.1 * self.config_field.velocity
        
        # Run simulation
        self.flow_field.run(self.sample_interval, control)
        
        # Get new state
        new_state = self.flow_field.obs.copy()[2:2 + self.n_sensors * 2]
        self.fifo_states.append(new_state)
        
        # Compute observation
        if len(self.fifo_states) >= 2:
            state_derivative = self.fifo_states[-1] - self.fifo_states[-2]
        else:
            state_derivative = np.zeros_like(new_state)
        
        obs = np.concatenate([
            self._normalize_state(new_state),
            self._normalize_state(state_derivative)
        ]).astype(self.dtype)
        
        # Compute reward
        reward = self._compute_reward(new_state, action)
        
        # Check termination
        self.current_step += 1
        terminated = False  # CFD simulations typically don't have natural termination
        truncated = self.current_step >= self.max_steps
        
        # Info
        info = {
            'step': self.current_step,
            'cd': float(self.flow_field.obs[0]),
            'cl': float(self.flow_field.obs[1]),
            'reward_cd': float(self.reward_cd),
            'reward_cl': float(self.reward_cl),
            'reward_sim': float(self.reward_sim),
        }
        
        return obs, reward, terminated, truncated, info
    
    def render(self):
        """Render the environment (not implemented)."""
        pass
    
    def close(self):
        """Clean up resources."""
        if hasattr(self, 'flow_field'):
            del self.flow_field