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 matplotlib.pyplot as plt
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 = 1200
FIFO_LEN = 120
CONV_LEN = 60
MAX_STEPS = 360
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, 8), 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 = 30
        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 / 2)
        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()
            self.target_states = np.vstack((self.target_states, new_state))

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

        for i in range(FIFO_LEN):
            self.flow_field.run(SAMPLE_INTERVAL, np.zeros(8, dtype=DATA_TYPE))
            self.fifo_states.append(self.flow_field.obs.copy())

        temp_states = np.array(self.fifo_states)
        self.force_norm_fact = 6 * np.max(np.abs(temp_states[:, 8:16]))
        for i in range(6):
            self.sens_deviation[i] = np.mean(temp_states[:, i+2])
            self.sens_norm_fact[i] = 5 * np.max(np.abs(temp_states[:, i+2] - self.sens_deviation[i]))
            self.target_states[:, i+2] = (self.target_states[:, i+2] - 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(8, dtype=DATA_TYPE)
                temp[5:8] = 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())

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

            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)

            similarities = 0.0
            target_seq = self.target_states[:, 2]
            state_seq = (states[:, 2] - self.sens_deviation[0]) / self.sens_norm_fact[0]
            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+2]
                state_seq = (states[:, i+2] - self.sens_deviation[i]) / self.sens_norm_fact[i]
                similarities += calc_sim(target_seq, state_seq, lag) / 6

            reward_sim = similarities

            target_seq = self.target_states[:, 0]
            state_seq = states[:, 8] + states[:, 10] + states[:, 12] + states[:, 14]
            ave_drag = np.average(state_seq)
            lag = calc_lag(target_seq, state_seq)
            similarities += calc_sim(target_seq, state_seq, lag) / 2
            target_seq = self.target_states[:, 1]
            state_seq = states[:, 9] + states[:, 11] + states[:, 13] + states[:, 15]
            similarities += calc_sim(target_seq, state_seq, lag) / 2
            reward_force = similarities

            reward_cd = np.exp(-np.abs((cd-ave_drag) * 2))
            reward_cl = np.exp(-np.abs(cl * 8))

            reward = np.minimum(0.0 * reward_cd + 0.0 * reward_cl + 0.4 * reward_force + 0.6 * reward_sim, 1.0)
            # barrier.wait()
            result_queue.put((np.hstack([forces[2:8], 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"):
        NX = self.flow_field.FIELD_SHAPE[0]
        NY = self.flow_field.FIELD_SHAPE[1]
        self.flow_field.get_ddf()
        ddf_plot = self.flow_field.ddf.copy().reshape((9, NY, NX)).transpose(2, 1, 0)
        ux = (ddf_plot[:, :, 1] + ddf_plot[:, :, 5] + ddf_plot[:, :, 8] - ddf_plot[:, :, 3] - ddf_plot[:, :, 6] - ddf_plot[:, :, 7]) / U0
        uy = (ddf_plot[:, :, 2] + ddf_plot[:, :, 5] + ddf_plot[:, :, 6] - ddf_plot[:, :, 4] - ddf_plot[:, :, 7] - ddf_plot[:, :, 8]) / U0
        speed = np.sqrt(ux**2 + uy**2)
        plt.figure(figsize=(10, 5))
        plt.imshow(speed.T, origin='lower', cmap='viridis', extent=[0, NX, 0, NY])
        plt.colorbar(label='Speed')
        plt.title('Scalar Velocity Field')
        plt.xlabel('X')
        plt.ylabel('Y')
        plt.tight_layout()
        plt.show()

    def save_field(self, filename):
        NX = self.flow_field.FIELD_SHAPE[0]
        NY = self.flow_field.FIELD_SHAPE[1]
        self.flow_field.get_ddf()
        ddf_plot = self.flow_field.ddf.copy().reshape((9, NY, NX)).transpose(2, 1, 0)
        flag_plot = self.flow_field.flag.copy().reshape((NY, NX)).transpose(1, 0)
        ux = (ddf_plot[:, :, 1] + ddf_plot[:, :, 5] + ddf_plot[:, :, 8] - ddf_plot[:, :, 3] - ddf_plot[:, :, 6] - ddf_plot[:, :, 7]) / U0
        uy = (ddf_plot[:, :, 2] + ddf_plot[:, :, 5] + ddf_plot[:, :, 6] - ddf_plot[:, :, 4] - ddf_plot[:, :, 7] - ddf_plot[:, :, 8]) / U0
        with open(os.path.join(parent_dir, "output", filename), "w") as f:
            f.write("Title= \"LBM 2D\"\r\n")
            f.write("VARIABLES= \"X\",\"Y\",\"flag\",\"U\",\"V\",\r\n")
            f.write(f"ZONE T= \"BOX\",I= {NX},J= {NY},F=POINT\r\n")
            for j in range(NY):
                for i in range(NX):
                    f.write(f"{i},{j},{flag_plot[i, j]},{ux[i, j]},{uy[i, j]}\r\n")

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