import numpy as np
import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec

from ..structs import RadiationProfiles, Parameters
from ..cosmo import T_adiab


#### Mhalo(z) from simulation from HM paper
z_array_1 = np.array([10.08083140877598, 9.705542725173212, 9.301385681293302, 8.897228637413393, 8.348729792147806, 7.684757505773671, 7.107390300230946, 6.5011547344110845, 5.92378752886836, 5.08660508083141])

Mh_z_1 = np.array([94706100.0077222, 126579158.66671978, 163154262.77379718, 222053024.50155312, 324945871.5591836, 530163041.15627587, 790019471.240894, 1243048829.5695167, 1991648063.8308542, 3689173954.4382358])

z_array_2 = np.array([12.361431870669746, 11.870669745958429, 11.379907621247115, 10.744803695150113, 9.821016166281755, 8.897228637413393, 8.146651270207853, 7.453810623556583, 6.732101616628176, 5.981524249422634, 5.057736720554274])

Mh_z_2 = np.array([98203277.90631257, 136100044.66105506, 195586368.30395073, 302214265.51783735, 570040236.7896342, 1114920997.0080914, 1852322215.6741183, 3191074972.923546, 5597981979.123267, 10000000000, 21414535638.539528])

z_array_3 = np.array([14.988452655889146, 14.642032332563511, 13.833718244803697, 12.70785219399538, 11.986143187066977, 11.23556581986143, 10.600461893764434, 9.79214780600462, 8.926096997690532, 8.146651270207853, 7.511547344110854, 6.876443418013858, 6.241339491916859, 5.6928406466512715, 5.288683602771364, 5.057736720554274])

Mh_z_3 = np.array([172274291.4769941, 218063348.75063255, 368917395.4438236, 775825016.8566794, 1243048829.5695167, 2102977594.5461233, 3493872774.7491226, 6129168695.9257145, 11987818459.583773, 21029775945.46132, 35577964903.39488, 61291686959.25715, 109488896512.76825, 178635801924.5735, 266193134612.6126, 343109759067.9875])



def plot_1D_profiles(parameters: Parameters, profiles: RadiationProfiles, mass_index: int, redshifts: list, alphas: list, label: str = None, figsize: tuple = (14, 9), fontsize: int = None, labelsize: int = None, cmap: str = 'plasma', **gridspec_kw) -> None:
    """Plot 1D radiation profiles for selected masses, redshifts and alphas.

    The figure has two rows:
    - Top row: halo mass accretion history (left) and legend (right).
    - Bottom row: Lyman-alpha flux, kinetic temperature, and ionisation fraction profiles.

    In the profile panels **colour encodes redshift** (blue → red, high z → low z)
    and **linestyle encodes alpha** (solid / dashed / dotted / dash-dot).

    In the MAR panel the Behroozi+20 simulation tracks are solid black curves;
    the analytical exponential MAR model is shown as dashed curves, one per
    alpha value, using the same colour palette as the profile panels' alpha
    linestyle legend (tab10).

    Args:
        parameters (Parameters): Simulation parameters used for axis labels and lookups.
        profiles (RadiationProfiles): Radiation profiles to plot.
        mass_index (int): Index selecting the mass bin to visualize.
        redshifts (list): Sequence of redshifts to plot (closest match is used).
        alphas (list): Sequence of alpha values to plot (closest match is used).
        label (str, optional): Optional suptitle for the figure.
        figsize (tuple, optional): Figure size as (width, height) in inches. Default is (14, 9).
        fontsize (int, optional): Font size for axis labels and legends.
            When ``None`` (default) the current ``rcParams`` values are used.
        labelsize (int, optional): Font size for tick labels.
            When ``None`` (default) the current ``rcParams`` values are used.
        cmap (str, optional): Colormap name used to colour redshift curves in the
            profile panels and scatter points in the MAR panel.  Default is
            ``'plasma'``.
        **gridspec_kw: Extra keyword arguments forwarded to :class:`matplotlib.gridspec.GridSpec`
            (e.g. ``hspace=0.4``, ``wspace=0.3``).  When ``hspace`` or ``wspace`` is supplied,
            ``constrained_layout`` is automatically disabled so the manual spacing takes effect.
    """
    import matplotlib.lines as mlines

    use_constrained_layout = 'hspace' not in gridspec_kw and 'wspace' not in gridspec_kw
    fig = plt.figure(figsize=figsize, constrained_layout=use_constrained_layout)
    fig.suptitle(label)

    gs = gridspec.GridSpec(2, 3, figure=fig, **gridspec_kw)
    ax_mar    = fig.add_subplot(gs[0, 0])
    ax_legend = fig.add_subplot(gs[0, 1:])
    ax_lyal   = fig.add_subplot(gs[1, 0])
    ax_temp   = fig.add_subplot(gs[1, 1])
    ax_xhii   = fig.add_subplot(gs[1, 2])

    co_radial_grid = profiles.r_grid_cell[:]
    r_lyal_phys    = profiles.r_lyal[:]
    zz             = profiles.z_history[:]

    n_z = len(redshifts)
    n_a = len(alphas)

    # Colour encodes redshift — sampled from the requested colormap
    _cmap = plt.get_cmap(cmap)
    colors_z = [_cmap(0.15 + 0.7 * i / max(n_z - 1, 1)) for i in range(n_z)]
    # Linestyle encodes alpha — shared between profile panels and MAR analytical curves
    ls_cycle = ['-', '--', ':', '-.']
    # Colours for the analytical MAR curves: C0, C1, C2, …
    colors_a = [f'C{j}' for j in range(n_a)]

    actual_z_vals    = []
    actual_alpha_vals = []

    for i, zi in enumerate(redshifts):
        ind_z = np.argmin(np.abs(zz - zi))
        z_val = zz[ind_z]
        actual_z_vals.append(z_val)

        for j, alpha_j in enumerate(alphas):
            ind_alpha = np.argmin(np.abs(parameters.solver.halo_mass_accretion_alpha - alpha_j))
            alpha_val = parameters.solver.halo_mass_accretion_alpha[ind_alpha]
            if i == 0:
                actual_alpha_vals.append(alpha_val)

            Mh_i = profiles.halo_mass_bins[mass_index, ind_alpha, ind_z]
            T_adiab_z    = T_adiab(z_val, parameters)
            Temp_profile = profiles.rho_heat[:, mass_index, ind_alpha, ind_z] + T_adiab_z

            x_HII_profile = np.zeros(len(co_radial_grid))
            x_HII_profile[co_radial_grid < profiles.R_bubble[mass_index, ind_alpha, ind_z]] = 1

            lyal_profile = profiles.rho_alpha[:, mass_index, ind_alpha, ind_z]

            color = colors_z[i]
            ls    = ls_cycle[j % len(ls_cycle)]

            ax_mar.scatter(z_val, Mh_i / 0.68, s=150, marker='*', color=color)
            ax_lyal.loglog(r_lyal_phys * (1 + z_val) / 0.68, lyal_profile, lw=1.7, color=color, ls=ls)
            ax_temp.loglog(co_radial_grid / 0.68, Temp_profile, lw=1.7, color=color, ls=ls)
            ax_xhii.semilogx(co_radial_grid / 0.68, x_HII_profile, lw=1.7, color=color, ls=ls)

    # MAR panel: Behroozi+20 simulation tracks — faint thick gray solid curves
    ax_mar.semilogy(z_array_1, Mh_z_1 / 0.68, color='gray', ls='-', lw=3, alpha=0.7)
    ax_mar.semilogy(z_array_2, Mh_z_2 / 0.68, color='gray', ls='-', lw=3, alpha=0.7)
    ax_mar.semilogy(z_array_3, Mh_z_3 / 0.68, color='gray', ls='-', lw=3, alpha=0.7)

    # MAR panel: analytical exponential MAR — one coloured curve per alpha,
    # linestyle matches the profile panels (ls_cycle[j])
    for j, alpha_j in enumerate(alphas):
        ind_alpha = np.argmin(np.abs(parameters.solver.halo_mass_accretion_alpha - alpha_j))
        ax_mar.semilogy(
            zz, profiles.halo_mass_bins[mass_index, ind_alpha, :] / 0.68,
            color=colors_a[j], ls=ls_cycle[j % len(ls_cycle)], lw=2,
        )

    # Build kwargs dicts
    label_kw  = {'fontsize': fontsize} if fontsize is not None else {}
    tick_kw   = {'labelsize': labelsize} if labelsize is not None else {}
    legend_kw = {'fontsize': fontsize} if fontsize is not None else {}

    # Style the MAR panel
    ax_mar.set_xlim(15, 5)
    ax_mar.set_ylim(1.5e8, 8e12)
    ax_mar.set_xlabel('z', **label_kw)
    ax_mar.set_ylabel(r'$M_h$ [$M_{\odot}$]', **label_kw)
    ax_mar.tick_params(axis='both', **tick_kw)

    # Legend panel: two stacked legends built from proxy artists
    ax_legend.axis('off')

    # Top legend: MAR model (Behroozi + one entry per alpha)
    mar_handles = [mlines.Line2D([], [], color='gray', ls='-', lw=3, alpha=0.7, label='Simulation (Behroozi+2020)')]
    for j, (alpha_val, ca) in enumerate(zip(actual_alpha_vals, colors_a)):
        mar_handles.append(
            mlines.Line2D([], [], color=ca, ls=ls_cycle[j % len(ls_cycle)], lw=2,
                          label=f'Analytical MAR  $\\alpha={alpha_val:.2f}$')
        )
    leg_mar = ax_legend.legend(handles=mar_handles, loc='upper center',
                               frameon=True, title='Halo accretion model', **legend_kw)
    ax_legend.add_artist(leg_mar)

    # Bottom legend: two-column layout — redshifts (left) | alphas (right)
    # Build each column separately, then interleave so ncol=2 fills them cleanly.
    _empty = mlines.Line2D([], [], color='none', label='')
    col_z = [mlines.Line2D([], [], color=cz, ls='-', lw=2, label=f'$z \\sim {z_val:.1f}$')
             for z_val, cz in zip(actual_z_vals, colors_z)]
    col_a = [mlines.Line2D([], [], color='k', ls=ls_cycle[j % len(ls_cycle)], lw=2,
                           label=f'$\\alpha = {alpha_val:.2f}$')
             for j, alpha_val in enumerate(actual_alpha_vals)]
    # Pad the shorter column with invisible entries so rows line up.
    # matplotlib ncol=2 fills the left column top-to-bottom first, then right.
    n = max(len(col_z), len(col_a))
    col_z += [_empty] * (n - len(col_z))
    col_a += [_empty] * (n - len(col_a))
    profile_handles = col_z + col_a
    ax_legend.legend(handles=profile_handles, ncol=2, loc='lower center',
                     frameon=True, title='colour = z  |  style = α', **legend_kw)

    # Style the profile panels
    ax_lyal.set_xlim(2e-1, 1e3)
    ax_lyal.set_ylim(2e-17, 1e-5)
    ax_lyal.set_xlabel('r [cMpc]', **label_kw)
    ax_lyal.set_ylabel(r'$\rho_{\alpha}$ [$\mathrm{pcm}^{-2}\, \mathrm{s}^{-1} \, \mathrm{Hz}^{-1}$]', **label_kw)
    ax_lyal.tick_params(axis='both', **tick_kw)

    ax_temp.set_xlim(2e-2, 1e2)
    ax_temp.set_ylim(0.8, 5e6)
    ax_temp.set_xlabel('r [cMpc]', **label_kw)
    ax_temp.set_ylabel(r'$T_k$ [K]', **label_kw)
    ax_temp.tick_params(axis='both', **tick_kw)

    ax_xhii.set_xlim(2e-2, 1e2)
    ax_xhii.set_ylim(0, 1.2)
    ax_xhii.set_xlabel('r [cMpc]', **label_kw)
    ax_xhii.set_ylabel(r'$x_{\mathrm{HII}}$', **label_kw)
    ax_xhii.tick_params(axis='both', **tick_kw)

    fig.show()


def plot_profile_alpha_dependence(ax: plt.axes, profiles: RadiationProfiles, quantity: str, mass_index: int, redshift_index: int, alphas: list, colors: list) -> None:

    # since these are hdf5 datasets, we need to copy it to a numpy array first
    co_radial_grid = profiles.r_grid_cell[:]
    r_lyal_phys = profiles.r_lyal[:]
    zz = profiles.z_history[:]

    z_val = zz[redshift_index]
    actual_alpha_list = []

    print(f"Plotting profiles at {redshift_index=}, {z_val=} => M_h = {profiles.halo_mass_bins[mass_index, 0, redshift_index]:.2e}")

    for j, alpha_j in enumerate(alphas):
        # the user specifies the redshifts and alpha values - here we find the index lying closest to these values in the profile

        ind_alpha = np.argmin(np.abs(profiles.parameters.solver.halo_mass_accretion_alpha - alpha_j))
        alpha_val = profiles.parameters.solver.halo_mass_accretion_alpha[ind_alpha]

        actual_alpha_list.append(alpha_val)

        # some quantities are required to plot sensible profiles
        T_adiab_z = T_adiab(z_val, profiles.parameters)
        Temp_profile = profiles.rho_heat[:, mass_index, ind_alpha, redshift_index] + T_adiab_z

        x_HII_profile = np.zeros((len(co_radial_grid)))
        x_HII_profile[np.where(co_radial_grid < profiles.R_bubble[mass_index, ind_alpha, redshift_index])] = 1

        lyal_profile = profiles.rho_alpha[:, mass_index, ind_alpha, redshift_index]  # *1.81e11/(1+zzi)

        color = colors[j]


        if quantity == 'lyal':
            ax.loglog(r_lyal_phys * (1 + z_val) / 0.68, lyal_profile, lw=1.7, color=color, alpha=0.8)
            ax.set_xlabel('r [cMpc]')
            ax.set_ylabel(r'$\rho_{\alpha}$ [$\mathrm{pcm}^{-2}\, \mathrm{s}^{-1} \, \mathrm{Hz}^{-1}$]')

        elif quantity == 'temp':
            ax.loglog(co_radial_grid / 0.68, Temp_profile, lw=1.7, color=color, alpha=0.8)

            ax.set_xlabel('r [cMpc]')
            ax.set_ylabel(r'$\rho_{h}$ [K]')


        elif quantity == 'xHII':
            ax.semilogx(co_radial_grid / 0.68, x_HII_profile, lw=1.7, color=color, alpha=0.8)
            ax.set_xlabel('r [cMpc]')
            ax.tick_params(axis="both")
            ax.set_ylabel(r'$x_{\mathrm{HII}}$')

        else:
            raise ValueError(f"Unknown quantity: {quantity}")
