import pickle
from dataclasses import dataclass
from enum import Enum
from pathlib import Path
from typing import List

import numpy as np
from matplotlib import pyplot as plt
from matplotlib.axes import Axes
from matplotlib.colors import LogNorm
from matplotlib.figure import Figure

from cic import cic_from_radius
from halo_mass_profile import halo_mass_profile
from paths import auriga_dir, richings_dir
from readfiles import read_file, read_halo_file
from utils import read_swift_config


class Mode(Enum):
    richings = 1
    auriga6 = 2


mode = Mode.richings


def dir_name_to_parameter(dir_name: str):
    return map(int, dir_name.lstrip("auriga6_halo").lstrip("richings21_").split("_"))


def levelmax_to_softening_length(levelmax: int) -> float:
    box_size = 100
    return box_size / 30 / 2 ** levelmax


fig1: Figure = plt.figure(figsize=(9, 6))
ax1: Axes = fig1.gca()
fig2: Figure = plt.figure(figsize=(9, 6))
ax2: Axes = fig2.gca()

for ax in [ax1, ax2]:
    ax.set_xlabel(r'R [Mpc]$')
ax1.set_ylabel(r'M [$10^{10} \mathrm{M}_\odot$]')
ax2.set_ylabel("density [$\\frac{10^{10} \\mathrm{M}_\\odot}{Mpc^3}$]")

part_numbers = []

reference_file = Path("auriga_reference.pickle")


@dataclass
class Result:
    title: str
    rho: np.ndarray


images = []
vmin = np.Inf
vmax = -np.Inf
root_dir = auriga_dir if mode == Mode.auriga6 else richings_dir
dirs = [d for d in root_dir.glob("*") if d.is_dir() and "bak" not in d.name]
for i, dir in enumerate(sorted(dirs)):
    is_by_adrian = "arj" in dir.name
    print(dir.name)

    if not is_by_adrian:
        levelmin, levelmin_TF, levelmax = dir_name_to_parameter(dir.name)
        print(levelmin, levelmin_TF, levelmax)
        # if levelmax != 12:
        #     continue
    Xc_adrian = 56.50153741810241
    Yc_adrian = 49.40761085700951
    Zc_adrian = 49.634393647291695
    Xc = 58.25576087992683
    Yc = 51.34632916228137
    Zc = 51.68749302578122

    input_file = dir / "output_0007.hdf5"
    if mode == Mode.richings:
        input_file = dir / "output_0004.hdf5"
    if is_by_adrian:
        input_file = dir / "output_0000.hdf5"
        softening_length = None
    else:
        swift_conf = read_swift_config(dir)
        softening_length = swift_conf["Gravity"]["comoving_DM_softening"]
        assert softening_length == swift_conf["Gravity"]["max_physical_DM_softening"]
        ideal_softening_length = levelmax_to_softening_length(levelmax)
        # if not np.isclose(softening_length, levelmax_to_softening_length(levelmax)):
        #     raise ValueError(f"softening length for levelmax {levelmax} should be {ideal_softening_length} "
        #                      f"but is {softening_length}")
    print(input_file)
    df, particles_meta = read_file(input_file)
    df_halos = read_halo_file(input_file.with_name("fof_" + input_file.name))
    # halos = read_velo_halos(dir, veloname="velo_out")
    # particles_in_halo = df.loc[df["FOFGroupIDs"] == 3]

    halo_id = 1
    while True:
        particles_in_halo = df.loc[df["FOFGroupIDs"] == halo_id]
        if len(particles_in_halo) > 1:
            break
        halo_id += 1

    halo = df_halos.loc[halo_id]
    part_numbers.append(len(df) * particles_meta.particle_mass)
    # halo = halos.loc[1]
    log_radial_bins, bin_masses, bin_densities, center = halo_mass_profile(
        df, halo, particles_meta, plot=False, num_bins=100,
        vmin=0.002, vmax=6.5
    )
    if is_by_adrian:
        with reference_file.open("wb") as f:
            pickle.dump([log_radial_bins, bin_masses, bin_densities], f)
    ax1.loglog(log_radial_bins[:-1], bin_masses, label=str(dir.name), c=f"C{i}")

    ax2.loglog(log_radial_bins[:-1], bin_densities, label=str(dir.name), c=f"C{i}")

    if reference_file.exists() and not is_by_adrian and mode == Mode.auriga6:
        with reference_file.open("rb") as f:
            data: List[np.ndarray] = pickle.load(f)
            ref_log_radial_bins, ref_bin_masses, ref_bin_densities = data
        mass_deviation: np.ndarray = np.abs(bin_masses - ref_bin_masses)
        density_deviation: np.ndarray = np.abs(bin_densities - ref_bin_densities)
        ax1.loglog(log_radial_bins[:-1], mass_deviation, c=f"C{i}",
                   linestyle="dotted")

        ax2.loglog(log_radial_bins[:-1], density_deviation, c=f"C{i}",
                   linestyle="dotted")
        accuracy = mass_deviation / ref_bin_masses
        print(accuracy)
        print("mean accuracy", accuracy.mean())

    if softening_length:
        for ax in [ax1, ax2]:
            ax.axvline(4 * softening_length, color=f"C{i}", linestyle="dotted")

    X, Y, Z = df.X.to_numpy(), df.Y.to_numpy(), df.Z.to_numpy()
    print()
    print(Yc - Yc_adrian)
    # shift: (-6, 0, -12)
    # if not is_by_adrian:
    #     xshift = Xc - Xc_adrian
    #     yshift = Yc - Yc_adrian
    #     zshift = Zc - Zc_adrian
    #     print("shift", xshift, yshift, zshift)

    X -= center[0]
    Y -= center[1]
    Z -= center[2]

    rho, extent = cic_from_radius(X, Z, 500, 0, 0, 5, periodic=False)

    vmin = min(vmin, rho.min())
    vmax = max(vmax, rho.max())

    images.append(Result(
        rho=rho,
        title=str(dir.name)
    ))

    # plot_cic(
    #     rho, extent,
    #     title=str(dir.name)
    # )
ax1.legend()
ax2.legend()

# fig3: Figure = plt.figure(figsize=(9, 9))
# axes: List[Axes] = fig3.subplots(3, 3, sharex=True, sharey=True).flatten()
fig3: Figure = plt.figure(figsize=(6, 3))
axes: List[Axes] = fig3.subplots(1, 2, sharex=True, sharey=True).flatten()

for result, ax in zip(images, axes):
    data = 1.1 + result.rho
    vmin_scaled = 1.1 + vmin
    vmax_scaled = 1.1 + vmax
    img = ax.imshow(data.T, norm=LogNorm(vmin=vmin_scaled, vmax=vmax_scaled), extent=extent,
                    origin="lower")
    ax.set_title(result.title)

fig3.tight_layout()
fig3.subplots_adjust(right=0.825)
cbar_ax = fig3.add_axes([0.85, 0.15, 0.05, 0.7])
fig3.colorbar(img, cax=cbar_ax)

fig1.savefig(Path(f"~/tmp/auriga1_{8}.pdf").expanduser())
fig2.savefig(Path(f"~/tmp/auriga2_{8}.pdf").expanduser())
fig3.savefig(Path("~/tmp/auriga3.pdf").expanduser())

plt.show()
print(part_numbers)