import gymnasium as gym
import numpy as np
from gymnasium import spaces
import ctypes
from collections import deque
from typing import Tuple
import sys
import os
import threading
from concurrent.futures import ThreadPoolExecutor
from concurrent.futures import ProcessPoolExecutor
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 = 12, 3
U0 = config_field.velocity
T0 = 1000
SAMPLE_INTERVAL = 800
FIFO_LEN = 120
CONV_LEN = 60
MAX_STEPS = 640
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):
    """Custom Environment that follows gym interface."""

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

    def __init__(self, device_id=0):
        super().__init__()
        self.action_space = spaces.Box(low=-1, high=1, shape=(A_DIM,), dtype=DATA_TYPE)
        self.observation_space = spaces.Box(
            low=-1, high=1, shape=(S_DIM,), dtype=DATA_TYPE
        )
        self.fifo_states = deque(maxlen=FIFO_LEN)
        self.target_states = np.empty((0, 6), dtype=DATA_TYPE)
        self.force_norm_fact = 1.0
        self.sens_norm_fact = np.ones(6, dtype=DATA_TYPE)
        self.sens_deviation = np.zeros(6, dtype=DATA_TYPE)

        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: Tuple[float, float, float] = (40 * L0, (NY - 1) / 2 + 2 * L0, 0)
        self.flow_field.add_sensor(center, L0 / 4)
        center: Tuple[float, float, float] = (40 * L0, (NY - 1) / 2, 0)
        self.flow_field.add_sensor(center, L0 / 4)
        center: Tuple[float, float, float] = (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 i 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))

        self.flow_field.apply_ddf()
        center: Tuple[float, float, float] = (30 * L0, (NY - 1) / 2, 0)
        self.flow_field.add_cylinder(center, L0 / 2)
        center: Tuple[float, float, float] = (31.3 * L0, (NY - 1) / 2 + 0.75 * L0, 0)
        self.flow_field.add_cylinder(center, L0 / 2)
        center: Tuple[float, float, float] = (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()

        for i 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]))
        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.target_states[:, i] = (self.target_states[:, i] - self.sens_deviation[i]) / self.sens_norm_fact[i]


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

        # barrier = threading.Barrier(2)
        result_queue = queue.Queue()

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

        def proc_data():
            states = np.array(self.fifo_states)
            forces = states[-1, 6:12] / self.force_norm_fact
            cd = (forces[0] + forces[2] + forces[4]) / 3
            cl = (forces[1] + forces[3] + forces[5]) / 3
            sens = (states[-1, 0:6] - self.sens_deviation) / self.sens_norm_fact

            similarities = 0.0

            def calc_lag(target, state):
                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))
                max_lag = lags[np.argmax(correlation)]
                return max_lag
                
            def calc_sim(target, state, lag):
                target_mean = np.mean(target)
                state_mean = np.mean(state)
                target_std = np.std(target)

                aligned_state = np.roll(state, lag)

                if lag >= 0:
                    seq_target = target[-CONV_LEN:]-target_mean
                    seq_state = aligned_state[-CONV_LEN:]-state_mean
                else:
                    seq_target = target[:CONV_LEN]-target_mean
                    seq_state = aligned_state[:CONV_LEN]-state_mean

                seq_diff = seq_target - seq_state
                sim_cor = 10*(np.corrcoef(seq_target, seq_state)[0, 1] - 1)
                sim_div = -np.abs((target_mean - state_mean) / target_std * 0.75)
                sim_amp = -np.abs(np.std(seq_diff) / target_std * 2)

                return np.exp((sim_cor + sim_div + sim_amp) / 3)

            id_sens = 0
            target_seq = self.target_states[:, id_sens]
            state_seq = (states[:, id_sens] - self.sens_deviation[id_sens]) / self.sens_norm_fact[id_sens]
            lag = calc_lag(target_seq, state_seq)
            similarities += calc_sim(target_seq, state_seq, lag) / 6

            for i in range(1, 6):
                target_seq = self.target_states[:, i]
                state_seq = (states[:, i] - self.sens_deviation[i]) / self.sens_norm_fact[i]
                similarities += calc_sim(target_seq, state_seq, lag) / 6

            reward_cd = np.exp(-np.abs(cd * 80))
            reward_cl = np.exp(-np.abs(cl * 20))
            # reward_sim = np.exp(2 * (similarities - 1))
            reward_sim = similarities
            reward = np.minimum(0.3 * reward_cd + 0.3 * reward_cl + 0.4 * reward_sim, 1.0)
            # barrier.wait()
            result_queue.put((np.hstack([forces, sens]), reward))

        run_flow_field(action)
        proc_data()
        observation, reward = result_queue.get()

        truncated = bool(np.any(observation > 1) or np.any(observation < -1))
        observation = np.clip(observation, -1, 1)
        # truncated = False
        return observation, float(reward), False, truncated, {}
    
    def reset(self, seed=None):
        self.flow_field.apply_ddf()
        return np.zeros(S_DIM, dtype=np.float32), {}
    
    def render(self, mode="human"):
        pass

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