import gymnasium as gym
import numpy as np
from gymnasium import spaces
from collections import deque
from typing import Tuple
import sys
import os
import queue

os.environ["OMP_NUM_THREADS"] = "1"
os.environ["MKL_NUM_THREADS"] = "1"

current_dir = os.path.dirname(os.path.abspath(__file__))
parent_dir = os.path.abspath(os.path.join(current_dir, os.pardir))
sys.path.append(parent_dir)
from CelerisLab import FlowField
from CelerisLab import utils

config_cuda = utils.load_cuda_config(
    os.path.join(parent_dir, "configs", "config_cuda.json")
)
config_field = utils.load_flow_field_config(
    os.path.join(parent_dir, "configs", "config_flowfield.json")
)

S_DIM, A_DIM = 2, 3
U0 = config_field.velocity
SAMPLE_INTERVAL = 800
FIFO_LEN = 150
CONV_LEN = 36
if config_field.data_type == "FP32":
    DATA_TYPE = np.float32
else:
    raise ValueError(f"Unsupported data type {config_field.data_type}.")


class CustomEnv(gym.Env):
    """DANTE-only environment with deterministic reset semantics and numeric-failure surfacing."""

    metadata = {"render_modes": ["human"], "render_fps": 1000 / SAMPLE_INTERVAL}

    FAILURE_NONE = 0
    FAILURE_NUMERIC = 1
    FAILURE_NONFINITE_OBS = 2
    FAILURE_OBS_OOB = 3

    def __init__(
        self,
        device_id: int = 0,
        obs_fail_bound: float = 2.0,
        obs_clip_bound: float = 3.0,
        reward_weights: Tuple[float, float, float] = (0.3, 0.3, 0.4),
    ):
        super().__init__()
        self.action_space = spaces.Box(low=-1, high=1, shape=(A_DIM,), dtype=DATA_TYPE)
        self.observation_space = spaces.Box(
            low=-obs_clip_bound,
            high=obs_clip_bound,
            shape=(S_DIM,),
            dtype=DATA_TYPE,
        )

        self.obs_fail_bound = float(obs_fail_bound)
        self.obs_clip_bound = float(obs_clip_bound)
        rw = np.asarray(reward_weights, dtype=DATA_TYPE)
        if rw.size != 3:
            raise ValueError("reward_weights must have length 3: (cd, cl, sim)")
        rw_sum = float(np.sum(rw))
        if rw_sum <= 0:
            raise ValueError("reward_weights sum must be positive")
        self.reward_weights = rw / rw_sum

        self.fifo_states = deque(maxlen=FIFO_LEN)
        self.target_states = np.empty((0, 6), dtype=DATA_TYPE)
        self.force_norm_fact = 1.0
        self.torque_norm_fact = 1.0
        self.sens_norm_fact = np.ones(6, dtype=DATA_TYPE)
        self.sens_deviation = np.zeros(6, dtype=DATA_TYPE)

        self.reward_cd = 0.0
        self.reward_cl = 0.0
        self.reward_sim = 0.0
        self.current_step = 0

        self.flow_field = FlowField(config_field, config_cuda, device_id)
        L0 = 20
        u0 = config_field.velocity
        nx = self.flow_field.FIELD_SHAPE[0]
        ny = self.flow_field.FIELD_SHAPE[1]

        center: Tuple[float, float, float] = (10 * L0, (ny - 1) / 2, 0)
        self.flow_field.add_cylinder(center, L0)
        center = (40 * L0, (ny - 1) / 2 + 2 * L0, 0)
        self.flow_field.add_sensor(center, L0 / 4)
        center = (40 * L0, (ny - 1) / 2, 0)
        self.flow_field.add_sensor(center, L0 / 4)
        center = (40 * L0, (ny - 1) / 2 - 2 * L0, 0)
        self.flow_field.add_sensor(center, L0 / 4)
        self.flow_field.run(int(4 * nx / u0), np.zeros(4, dtype=DATA_TYPE))

        for _ in range(FIFO_LEN):
            self.flow_field.run(SAMPLE_INTERVAL, np.zeros(4, dtype=DATA_TYPE))
            new_state = self.flow_field.obs.copy()[2:8]
            self.target_states = np.vstack((self.target_states, new_state))

        center = (30 * L0, (ny - 1) / 2, 0)
        self.flow_field.add_cylinder(center, L0 / 2)
        center = (31.3 * L0, (ny - 1) / 2 + 0.75 * L0, 0)
        self.flow_field.add_cylinder(center, L0 / 2)
        center = (31.3 * L0, (ny - 1) / 2 - 0.75 * L0, 0)
        self.flow_field.add_cylinder(center, L0 / 2)
        self.flow_field.run(int(4 * nx / u0), np.zeros(7, dtype=DATA_TYPE))
        self.flow_field.get_ddf()
        self.flow_field.save_ddf()

        for _ in range(FIFO_LEN):
            self.flow_field.run(SAMPLE_INTERVAL, np.zeros(7, dtype=DATA_TYPE))
            self.fifo_states.append(self.flow_field.obs.copy()[2:14])

        temp_states = np.array(self.fifo_states)
        self.force_norm_fact = 6 * np.max(np.abs(temp_states[:, 6:12]))
        temp_torque = (
            -temp_states[:, 1]
            - temp_states[:, 2] * np.sqrt(3) / 2
            + temp_states[:, 3] / 2
            + temp_states[:, 4] * np.sqrt(3) / 2
            + temp_states[:, 5] / 2
        )
        self.torque_norm_fact = 10 * np.max(np.abs(temp_torque))

        for i in range(6):
            self.sens_deviation[i] = np.mean(temp_states[:, i])
            self.sens_norm_fact[i] = 5 * np.max(np.abs(temp_states[:, i] - self.sens_deviation[i]))

        self.flow_field.apply_ddf()

        for _ in range(FIFO_LEN):
            self.flow_field.run(
                SAMPLE_INTERVAL,
                np.array([0.0, 0.0, 0.0, 0.0, 0.0, -4 * u0, 4 * u0], dtype=DATA_TYPE),
            )
            self.fifo_states.append(self.flow_field.obs.copy()[2:14])

        self.save_states = self.fifo_states.copy()
        # self.flow_field.apply_ddf()
        self.flow_field.get_ddf()
        self.flow_field.save_ddf()

    def _calc_lag(self, target: np.ndarray, state: np.ndarray) -> int:
        target_mean = np.mean(target)
        state_mean = np.mean(state)
        correlation = np.correlate(target - target_mean, state - state_mean, "full")
        lags = np.arange(-len(target) + 1, len(target))
        return int(lags[np.argmax(correlation)])

    def _calc_dtw_sim(self, target: np.ndarray, state: np.ndarray) -> float:
        n = len(target)
        m = len(state)
        dtw_matrix = np.full((n + 1, m + 1), np.inf)
        dtw_matrix[0, 0] = 0
        for i in range(1, n + 1):
            for j in range(1, m + 1):
                cost = abs(target[i - 1] - state[j - 1])
                last_min = min(
                    dtw_matrix[i - 1, j],
                    dtw_matrix[i, j - 1],
                    dtw_matrix[i - 1, j - 1],
                )
                dtw_matrix[i, j] = cost + last_min
        return float(1 - (dtw_matrix[n, m] / len(target)))

    def _compute_obs_reward(self):
        states = np.array(self.fifo_states)
        forces = states[-1, 6:12] / self.force_norm_fact

        obs_drag = float((forces[0] + forces[2] + forces[4]) / 3)
        obs_lift = float((forces[1] + forces[3] + forces[5]) / 3)

        similarities = 0.0
        id_sens = 1
        target_seq = self.target_states[CONV_LEN : 2 * CONV_LEN, id_sens]
        state_seq = states[-CONV_LEN:, id_sens]
        lag = self._calc_lag(target_seq, state_seq)

        for i in range(0, 6):
            target_seq = np.roll(self.target_states[:, i], -lag)[CONV_LEN : 2 * CONV_LEN]
            state_seq = states[-CONV_LEN:, i]
            similarities += self._calc_dtw_sim(target_seq, state_seq) / 6

        self.reward_cd = float(np.exp(-np.abs(obs_drag * 20)))
        self.reward_cl = float(np.exp(-np.abs(obs_lift * 80)))
        self.reward_sim = float(np.exp(-10 * np.abs(similarities - 1)))
        reward = float(
            np.minimum(
                self.reward_weights[0] * self.reward_cd
                + self.reward_weights[1] * self.reward_cl
                + self.reward_weights[2] * self.reward_sim,
                1.0,
            )
        )
        observation = np.array([forces[0], forces[1]], dtype=DATA_TYPE)
        return observation, reward

    def step(self, action):
        assert self.action_space.contains(action), "%r (%s) invalid" % (
            action,
            type(action),
        )

        result_queue = queue.Queue()

        def run_flow_field(act):
            self.flow_field.context.push()
            u0 = config_field.velocity
            try:
                temp = np.zeros(7, dtype=DATA_TYPE)
                temp[4:7] = np.array((act * 8 + [0, -4, 4]) * u0, dtype=DATA_TYPE)
                self.flow_field.run(SAMPLE_INTERVAL, temp)
            finally:
                self.flow_field.context.pop()
            self.fifo_states.append(self.flow_field.obs.copy()[2:14])

        def proc_data():
            result_queue.put(self._compute_obs_reward())

        run_flow_field(action)

        if self.flow_field.has_numeric_error():
            self.current_step += 1
            obs = np.zeros(S_DIM, dtype=DATA_TYPE)
            info = {
                "failure_code": self.FAILURE_NUMERIC,
                "numeric_error": True,
                "raw_obs": obs.copy(),
            }
            return obs, 0.0, False, True, info

        proc_data()
        observation, reward = result_queue.get()
        raw_obs = np.asarray(observation, dtype=DATA_TYPE).copy()

        failure_code = self.FAILURE_NONE
        if not np.all(np.isfinite(observation)):
            failure_code = self.FAILURE_NONFINITE_OBS
        elif np.any(np.abs(observation) > self.obs_fail_bound):
            failure_code = self.FAILURE_OBS_OOB

        truncated = failure_code != self.FAILURE_NONE
        if truncated:
            reward = 0.0
            observation = np.zeros(S_DIM, dtype=DATA_TYPE)
        else:
            observation = np.clip(observation, -self.obs_clip_bound, self.obs_clip_bound)

        self.current_step += 1
        info = {
            "failure_code": int(failure_code),
            "numeric_error": False,
            "raw_obs": raw_obs,
        }
        return observation.astype(np.float32), float(reward), False, bool(truncated), info

    def reset(self, seed=None):
        self.flow_field.restore_ddf()
        self.flow_field.apply_ddf()
        self.fifo_states = self.save_states.copy()
        self.current_step = 0
        return np.zeros(S_DIM, dtype=np.float32), {}

    def close(self):
        self.flow_field.__del__()
