"""This module contains methods for data analysis"""

import os
from typing import Optional, List, Dict

import numpy as np
from h5py import File as H5File
import aptools.data_analysis.beam_diagnostics as bd

from wake_t.particles.particle_bunch import ParticleBunch


def analyze_bunch(
    bunch: ParticleBunch,
    n_slices: Optional[int] = 50,
    len_slice: Optional[float] = None,
) -> Dict:
    # perform analysis
    dist = bunch.prop_distance
    params_analysis = _get_distribution_parameters(
        bunch.x,
        bunch.y,
        bunch.xi,
        bunch.px,
        bunch.py,
        bunch.pz,
        bunch.q,
        n_slices,
        len_slice,
    )

    # store data
    bunch_params = _store_bunch_parameters_into_dict(dist, *params_analysis)
    return bunch_params


def analyze_bunch_list(
    bunch_list: List[ParticleBunch],
    n_slices: Optional[int] = 50,
    len_slice: Optional[float] = None,
) -> Dict:
    # preallocate arrays
    list_len = len(bunch_list)
    a_x = np.zeros(list_len)
    a_y = np.zeros(list_len)
    b_x = np.zeros(list_len)
    b_y = np.zeros(list_len)
    g_x = np.zeros(list_len)
    g_y = np.zeros(list_len)
    ene = np.zeros(list_len)
    dist = np.zeros(list_len)
    ene_sp = np.zeros(list_len)
    ene_sp_sl_avg = np.zeros(list_len)
    em_x = np.zeros(list_len)
    em_y = np.zeros(list_len)
    em_x_sl_avg = np.zeros(list_len)
    em_y_sl_avg = np.zeros(list_len)
    s_x = np.zeros(list_len)
    s_y = np.zeros(list_len)
    x_avg = np.zeros(list_len)
    y_avg = np.zeros(list_len)
    i_peak = np.zeros(list_len)
    theta_x = np.zeros(list_len)
    theta_y = np.zeros(list_len)
    s_z = np.zeros(list_len)
    q_tot = np.zeros(list_len)

    # perform analysis
    for i, bunch in enumerate(bunch_list):
        dist[i] = bunch.prop_distance
        params_analysis = _get_distribution_parameters(
            bunch.x,
            bunch.y,
            bunch.xi,
            bunch.px,
            bunch.py,
            bunch.pz,
            bunch.q,
            n_slices,
            len_slice,
        )
        (
            theta_x[i],
            theta_y[i],
            x_avg[i],
            y_avg[i],
            s_x[i],
            s_y[i],
            a_x[i],
            a_y[i],
            b_x[i],
            b_y[i],
            g_x[i],
            g_y[i],
            em_x[i],
            em_y[i],
            em_x_sl_avg[i],
            em_y_sl_avg[i],
            ene[i],
            ene_sp[i],
            ene_sp_sl_avg[i],
            i_peak[i],
            s_z[i],
            q_tot[i],
        ) = params_analysis

    # store into dictionary
    bunch_list_params = _store_bunch_parameters_into_dict(
        dist,
        theta_x,
        theta_y,
        x_avg,
        y_avg,
        s_x,
        s_y,
        a_x,
        a_y,
        b_x,
        b_y,
        g_x,
        g_y,
        em_x,
        em_y,
        em_x_sl_avg,
        em_y_sl_avg,
        ene,
        ene_sp,
        ene_sp_sl_avg,
        i_peak,
        s_z,
        q_tot,
    )
    return bunch_list_params


def save_parameters_to_file(bunch_params, folder_path, file_name):
    if not os.path.exists(folder_path):
        os.makedirs(folder_path)
    file_path = os.path.join(folder_path, file_name + ".h5")
    with H5File(file_path, "w") as h5_file:
        for param_name, param_data in bunch_params.items():
            h5_file.create_dataset(param_name, data=param_data)


def read_parameters_from_file(file_path):
    bunch_params = {}
    with H5File(file_path, "r") as h5_file:
        for param in h5_file.keys():
            bunch_params[param] = h5_file.get(param).value
    return bunch_params


def save_bunch_to_file(bunch, folder_path, file_name):
    if not os.path.exists(folder_path):
        os.makedirs(folder_path)
    file_path = os.path.join(folder_path, file_name + ".h5")
    with H5File(file_path, "w") as h5_file:
        h5_file.create_dataset("x", data=bunch.x)
        h5_file.create_dataset("y", data=bunch.y)
        h5_file.create_dataset("xi", data=bunch.xi)
        h5_file.create_dataset("px", data=bunch.px)
        h5_file.create_dataset("py", data=bunch.py)
        h5_file.create_dataset("pz", data=bunch.pz)
        h5_file.create_dataset("q", data=bunch.q)
        h5_file.attrs["prop_dist"] = bunch.prop_distance


def _get_distribution_parameters(x, y, z, px, py, pz, q, n_slices, len_slice):
    a_x, b_x, g_x = bd.twiss_parameters(x, px, pz, py, w=q)
    a_y, b_y, g_y = bd.twiss_parameters(y, py, pz, px, w=q)
    ene = bd.mean_energy(px, py, pz, w=q)
    ene_sp = bd.relative_rms_energy_spread(px, py, pz, w=q)
    ene_sp_slice_params = bd.relative_rms_slice_energy_spread(
        z, px, py, pz, w=q, n_slices=n_slices, len_slice=len_slice
    )
    ene_sp_sl, sl_w, sl_e, ene_sp_sl_avg = ene_sp_slice_params
    em_x = bd.normalized_transverse_rms_emittance(x, px, py, pz, w=q)
    em_y = bd.normalized_transverse_rms_emittance(y, py, px, pz, w=q)
    emitt_x_slice_params = bd.normalized_transverse_rms_slice_emittance(
        z, x, px, py, pz, w=q, n_slices=n_slices, len_slice=len_slice
    )
    em_x_sl, sl_w, sl_e, em_x_sl_avg = emitt_x_slice_params
    emitt_y_slice_params = bd.normalized_transverse_rms_slice_emittance(
        z, y, py, px, pz, w=q, n_slices=n_slices, len_slice=len_slice
    )
    em_y_sl, sl_w, sl_e, em_y_sl_avg = emitt_y_slice_params
    s_z = bd.rms_length(z, w=q)
    s_x = bd.rms_size(x, w=q)
    s_y = bd.rms_size(y, w=q)
    x_avg = np.average(x, weights=q)
    y_avg = np.average(y, weights=q)
    px_avg = np.average(px, weights=q)
    py_avg = np.average(py, weights=q)
    theta_x = px_avg / ene
    theta_y = py_avg / ene
    i_peak = bd.peak_current(z, q, n_slices=n_slices, len_slice=len_slice)
    q_tot = np.sum(q)
    return (
        theta_x,
        theta_y,
        x_avg,
        y_avg,
        s_x,
        s_y,
        a_x,
        a_y,
        b_x,
        b_y,
        g_x,
        g_y,
        em_x,
        em_y,
        em_x_sl_avg,
        em_y_sl_avg,
        ene,
        ene_sp,
        ene_sp_sl_avg,
        i_peak,
        s_z,
        q_tot,
    )


def _store_bunch_parameters_into_dict(
    dist,
    theta_x,
    theta_y,
    x_avg,
    y_avg,
    s_x,
    s_y,
    a_x,
    a_y,
    b_x,
    b_y,
    g_x,
    g_y,
    em_x,
    em_y,
    em_x_sl_avg,
    em_y_sl_avg,
    ene,
    ene_sp,
    ene_sp_sl_avg,
    i_peak,
    s_z,
    q_tot,
):
    params_dict = {
        "prop_dist": dist,
        "theta_x": theta_x,
        "theta_y": theta_y,
        "x_avg": x_avg,
        "y_avg": y_avg,
        "sigma_x": s_x,
        "sigma_y": s_y,
        "sigma_z": s_z,
        "alpha_x": a_x,
        "alpha_y": a_y,
        "beta_x": b_x,
        "beta_y": b_y,
        "gamma_x": g_x,
        "gamma_y": g_y,
        "emitt_x": em_x,
        "emitt_y": em_y,
        "avg_slice_emitt_x": em_x_sl_avg,
        "avg_slice_emitt_y": em_y_sl_avg,
        "avg_ene": ene,
        "rel_ene_spread": ene_sp,
        "avg_slice_rel_ene_spread": ene_sp_sl_avg,
        "i_peak": i_peak,
        "q_tot": q_tot,
    }
    return params_dict
