# 这个是imit圆柱的env，用于产生d1a3o14_250525_imit系列模型的目标流场
# 跟随模型，要匹配圆柱直径和采样频率
# 模型名中L代表目标直径，S代表SAMPLE_INTERVAL
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 = 14, 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, 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.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]

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

        center: Tuple[float, float, float] = (20 * L0, (NY - 1) / 2, 0)
        self.flow_field.add_cylinder(center, 1*L0) # 这里会随着圆柱尺寸变动，在模型名中体现
        center: Tuple[float, float, float] = (30 * L0, (NY - 1) / 2 + 2 * L0, 0)
        self.flow_field.add_sensor(center, L0 / 4)
        center: Tuple[float, float, float] = (30 * L0, (NY - 1) / 2, 0)
        self.flow_field.add_sensor(center, L0 / 4)
        center: Tuple[float, float, float] = (30 * 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()[0:8]
            self.target_states = np.vstack((self.target_states, new_state))

        def analyze_harmonics(states, n_harmonics):
            N, D = states.shape
            result = []
            for d in range(D):
                y = states[:, d]
                fft_coef = np.fft.rfft(y)
                freqs = np.fft.rfftfreq(N, d=1)
                amps = 2 * np.abs(fft_coef) / N
                phases = np.angle(fft_coef)
                idx = np.argsort(amps[1:])[::-1][:n_harmonics] + 1
                harmonics = {
                    'dc': np.real(fft_coef[0]) / N,
                    'amps': amps[idx],
                    'freqs': freqs[idx],
                    'phases': phases[idx]
                }
                result.append(harmonics)
            return result

        self.target_harmonics = analyze_harmonics(self.target_states, n_harmonics=5)

        self.flow_field.get_ddf()
        self.flow_field.save_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(4, dtype=DATA_TYPE)
                self.flow_field.run(SAMPLE_INTERVAL, temp)
            finally:
                self.flow_field.context.pop()
            # barrier.wait()

        run_flow_field(action)

        truncated = False
        observation = np.zeros(14, dtype=DATA_TYPE)
        self.current_step += 1
        # done = self.current_step >= MAX_STEPS
        done = False
        return observation, float(0), done, truncated, {}
    
    def reset(self, seed=None):
        self.flow_field.restore_ddf()
        self.flow_field.apply_ddf()
        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__()