# 这个是erase的env，用于训练和评估d1a3o12_250729_250326_erase系列模型
# 上游扰流圆柱，场景与Karman_cloak_standard一致，
# 但是目标是希望pinball后流场跟入口一致，即抹除扰流圆柱尾迹
# 模型名中D代表信号延迟
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 = 600
FIFO_LEN = 150
CONV_LEN = 36
MAX_STEPS = 500
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.reward_u = 0.0
        self.reward_v = 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]

        self.ddf_ave = np.zeros((NX, NY, 9), dtype=DATA_TYPE)
        self.ddf_ave_cont = 0

        # 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(2*NX/U0), np.zeros(3, dtype=DATA_TYPE))

        for i in range(FIFO_LEN):
            self.flow_field.run(SAMPLE_INTERVAL, np.zeros(3, dtype=DATA_TYPE))
            new_state = self.flow_field.obs.copy()[0:6]
            self.target_states = np.vstack((self.target_states, new_state))

        self.target_sensors = np.mean(self.target_states, axis=0)

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

        temp_states = np.array(self.fifo_states)
        self.force_norm_fact = 100 * np.max(np.abs(temp_states[:, 8:14]))

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

        self.flow_field.apply_ddf()

        for i in range(FIFO_LEN):
            self.flow_field.run(SAMPLE_INTERVAL, np.array([0.0, 0.0, 0.0, 0.0, 0.0, -8*U0, 8*U0], dtype=DATA_TYPE))
            self.fifo_states.append(self.flow_field.obs.copy()[0:14])
            
        self.save_states = self.fifo_states.copy()
        self.flow_field.apply_ddf()

    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,-8,8])*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()[0:14])

        def proc_data():
            states = np.array(self.fifo_states)
            forces = states[-1, 6:14] / self.force_norm_fact
            cd = (forces[2] + forces[4] + forces[6]) / 3
            cl = (forces[3] + forces[5] + forces[7]) / 3
            sens = (states[-1, 0:6] - self.sens_deviation) / self.sens_norm_fact
            target_sens = (self.target_sensors - 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):
                # 计算幅值差异权重
                target_std = np.std(target) if np.std(target) > 1e-8 else 1e-8
                state_std = np.std(state) if np.std(state) > 1e-8 else 1e-8
                amplitude_ratio = min(target_std, state_std) / max(target_std, state_std)
                
                # 计算均值差异
                mean_diff = abs(np.mean(target) - np.mean(state))
                max_scale = max(abs(np.mean(target)), abs(np.mean(state)), 1e-8)
                mean_similarity = 1 / (1 + mean_diff / max_scale * 10)
                
                # DTW计算
                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
                
                # 改进的归一化方法
                max_possible_cost = max(np.max(np.abs(target)), np.max(np.abs(state)), 1e-8)
                dtw_distance = dtw_matrix[n, m] / (len(target) * max_possible_cost)
                DTW_similarity = max(0, 1 - dtw_distance)
                
                # 综合相似度：形状相似度 * 幅值相似度 * 均值相似度
                total_similarity = 0.8 * DTW_similarity + 0.1 * amplitude_ratio + 0.1 * mean_similarity
                
                return total_similarity

            # id_sens = 1
            # target_seq = self.target_states[CONV_LEN:2*CONV_LEN, id_sens]
            target_seq = -states[CONV_LEN:2*CONV_LEN, 7]
            # state_seq = states[-CONV_LEN:, id_sens]
            state_seq = states[-CONV_LEN:, 9]
            lag = calc_lag(target_seq, state_seq)

            for i in range(0, 2):
                target_seq = -np.roll(states[:, i+6], -lag)[CONV_LEN:2*CONV_LEN]
                state_seq = states[-CONV_LEN:, i+8] + states[-CONV_LEN:, i+10] + states[-CONV_LEN:, i+12]
                similarities += calc_sim(target_seq, state_seq) / 2

            diff_u = (np.abs(sens[0] - target_sens[0]) + np.abs(sens[2] - target_sens[2]) + np.abs(sens[4] - target_sens[4]))/3
            diff_v = (np.abs(sens[1] - target_sens[1]) + np.abs(sens[3] - target_sens[3]) + np.abs(sens[5] - target_sens[5]))/3
            mean_u = np.mean(np.abs(states[:, 0] - self.target_sensors[0]) \
                             + np.abs(states[:, 2] - self.target_sensors[2]) \
                             + np.abs(states[:, 4] - self.target_sensors[4])) / self.sens_norm_fact[2]
            amp_v = np.std(states[:, 1] + states[:, 3] + states[:, 5]) / self.sens_norm_fact[3]
            self.reward_u = np.exp(-np.abs(diff_u * 40))
            self.reward_v = 0.7 * np.exp(-np.abs(amp_v * 20)) + 0.3 * np.exp(-np.abs(diff_v * 20))
            # self.reward_sim = np.exp(-50*np.abs(similarities - 1)**2)
            self.reward_sim = similarities
            reward = np.minimum(0.4 * self.reward_u + 0.4 * self.reward_v + 0.2 * self.reward_sim, 1.0)
            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)
        self.current_step += 1
        # done = self.current_step >= MAX_STEPS
        done = False
        return observation, float(reward), done, truncated, {}
    
    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 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 average_field(self, mode=["add", "save", "clear"], filename="average_field.dat"):
        NX = self.flow_field.FIELD_SHAPE[0]
        NY = self.flow_field.FIELD_SHAPE[1]
        self.flow_field.get_ddf()
        ddf_new = self.flow_field.ddf.copy().reshape((9, NY, NX)).transpose(2, 1, 0)
        if "add" in mode:
            self.ddf_ave = self.ddf_ave + ddf_new
            self.ddf_ave_cont += 1
        if "save" in mode:
            if self.ddf_ave_cont == 0:
                raise ValueError("No data to save. Please run 'add' mode first.")
            ux = (self.ddf_ave[:, :, 1] + self.ddf_ave[:, :, 5] + self.ddf_ave[:, :, 8] - self.ddf_ave[:, :, 3] - self.ddf_ave[:, :, 6] - self.ddf_ave[:, :, 7]) / U0 / self.ddf_ave_cont
            uy = (self.ddf_ave[:, :, 2] + self.ddf_ave[:, :, 5] + self.ddf_ave[:, :, 6] - self.ddf_ave[:, :, 4] - self.ddf_ave[:, :, 7] - self.ddf_ave[:, :, 8]) / U0 / self.ddf_ave_cont
            flag_plot = self.flow_field.flag.copy().reshape((NY, NX)).transpose(1, 0)
            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")
            print(f"Average field amount: {self.ddf_ave_cont}")
        if "clear" in mode:
            self.ddf_ave = np.zeros((NX, NY, 9), dtype=DATA_TYPE)
            self.ddf_ave_cont = 0

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