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

def white_to_alpha(r, g, b):
    """Add transparency preserving the color when blended with white.
    
    RGB color components are mapped to HLS. A new color with lower luminance
    (L_new < 0.5) and higher saturation (S_new = 1) and a transparency value
    A_new are then computed, in such a way that the new color produces the
    original color if blended with a white background.
    
    Parameters
    ----------
    r, g, b : float
        RGB components of the original color, in the [0, 1] range.

    Returns
    -------
    r, g, b, a : float
        RGBA components of the new color, with white component made transparent.
    """
    import colorsys
    
    h, l, s = colorsys.rgb_to_hls(r, g, b)
    if l > 0.5:
        l_new = s/(1+s)
    else:
        l_new = 1 - (1-l)/(1 + l*(s-1))
    a = (1-l)/(1-l_new)
    r_new, g_new, b_new = colorsys.hls_to_rgb(h, l_new, 1)
    return r_new, g_new, b_new, a


def cmap_lin_alpha(name, min_loc=0, reverse_alpha=False, register=True):
    """
    Set linear alpha (transparency) value in a colormap.
    
    A new colormap is returned, with unchanged RGB values, and transparency
    changing linearly from A=0 (fully transparent) at min_alpha_loc to A=1
    (fully opaque) at the range ends.
    
    Parameters
    ----------
    name : str
        Name of builtin (or registered) colormap.
    min_loc : float, default=0
        Location in the map of maximum transparency.
        If min_loc=0, transparency decreases linearly from the lower end to
        the upper end of the map.
        If min_loc=1, transparency increases linearly from the lower end to
        the upper end of the map.
        If 0<min_loc<1, colors are fully opaque at both ends and transparency
        increases linearly towards min_loc. 
    reverse_alpha : bool, default=False
        Whether alpha channel is reversed, so that minimum transparency
        occurs at min_loc and maximum transparency at range ends.
    register : bool, default=True
        Whether to register the new colormap.

    Returns
    -------
    cmap : matplotlib.colors.ListedColormap
        Colormap with linear alpha set.
    """
    # Fetch the cmap callable
    cmap = plt.get_cmap(name)
    # Get the colors out from the colormap LUT
    rgba = cmap(np.arange(cmap.N)) # N-by-4
    # Map the max transparency location from (0, 1) to index in LUT
    index_alpha_min = int(np.rint(np.clip(min_loc*cmap.N, 0, cmap.N)))
    # Build the linear alpha channel with minimum at index_alpha_min
    alpha = np.concatenate(
        (np.linspace(1, 0, index_alpha_min, endpoint=False),
         np.linspace(0, 1, cmap.N - index_alpha_min)))
    # Replace the old alpha channel
    rgba[:,-1] = alpha
    # Create a new Colormap object
    cmap_alpha = mpl.colors.ListedColormap(rgba, name=name + '_lin_alpha')
    if register:
        mpl.cm.register_cmap(name=name + '_lin_alpha', cmap=cmap_alpha)
    return cmap_alpha


def cmap_transp(name, register=True):
    """
    Make transparent the white component in a colormap.
    
    A new colormap is returned, with new RGBA components that produce the
    original RGB colors (with A=1) if blended with a white background.
    
    Parameters
    ----------
    name : str
        Name of builtin (or registered) colormap.
    register : bool, default=True
        Whether to register the new colormap.

    Returns
    -------
    cmap : matplotlib.colors.ListedColormap
        The new colormap.
    """
    # Fetch the cmap callable
    cmap = plt.get_cmap(name)
    # Get the colors out from the colormap LUT
    rgba = cmap(np.arange(cmap.N)) # N-by-4
    # Replace the old alpha channel
    rgba = [white_to_alpha(*color[:3]) for color in rgba]
    # Create a new Colormap object
    cmap_alpha = mpl.colors.ListedColormap(rgba, name=name + '_transp')
    if register:
        mpl.cm.register_cmap(name=name + '_transp', cmap=cmap_alpha)
    return cmap_alpha


def cmap_partial(cmap, vmin=0, vmax=1):
    """
    Make a new colormap from a section of an existing one.
    
    A new colormap is returned, with the same number of colors as the original
    one, obtained by linear interpolation in a sub-range of the existing
    colormap LUT.
    
    Parameters
    ----------
    name : matplotlib.Colormap
        The existing colormap a sub-section of which will become the new
        colormap.
    vmin : the lower end of the range in the LUT that will be mapped to the
        minimum value (0) in the new colormap
    vmax : the upper end of the range in the LUT that will be mapped to the
        maximum value (1) in the new colormap

    Returns
    -------
    partial : matplotlib.colors.LinearSegmentedColormap
        The new colormap.
    """
    partial = mpl.colors.LinearSegmentedColormap.from_list(
        cmap.name + '_partial',
        cmap(np.linspace(vmin, vmax, cmap.N)))
    return partial


def read_grayout(dirname, **kwargs):
    """
    Make a new colormap from a section of an existing one.
    
    A new colormap is returned, with the same number of colors as the original
    one, obtained by linear interpolation in a sub-range of the existing
    colormap LUT.
    
    Parameters
    ----------
    dirname : str
        The path to the directory containing the output files of GRAY.
    **kwargs : optional keyword arguments
        Not yet implemented. Filenames will be passed here to override the
        default (fort.*) values.

    Returns
    -------
    run : dict
        All the output data found in dirname, organized as a dictionary. The
        field names in the structured arrays are derived from the column
        headers in the output files.
        - 'name' : str
            The base name of the directory
        - 'bres' : numpy.ndarray
            Structured array containing the contours of the EC resonances.
        - 'flux' : numpy.ndarray
            Structured array conaining the contours of the poloidal flux.
        - 'prof' : numpy.ndarray
            Structured array containing the ECCD profiles.
        - 'ray' : numpy.ndarray
            Structured array describing the central ray of the beam.
        - 'rays' : numpy.ndarray
            Structured array describing the envelope of the beam.
        - 'totals' : numpy.ndarray
            Structured array conaining the integrated, 0D quantities.
    """
    import os
    import subprocess
    import io

    HEADER_GRAY = 21
    UNITS_GRAY = {
        'ray': 4,
        'totals': 7,
        'rays': 33,
        'prof': 48,
        'bres': 70,
        'flux': 71,
        }

    if not os.path.isdir(dirname):
        print(f"{dirname} not a directory. Skipping...")
        return None
    if dirname.endswith("/"): dirname = dirname[:-1]

    name = os.path.basename(dirname)
    run = dict(name=name)
    for label, unit in UNITS_GRAY.items():
        file = dirname + f"/fort.{unit}"
        if label == 'bres' or label == 'flux':
            # Replace blank lines with as many NaNs as the number of columns.
            # In case of success, file is a string with the file's content.
            # If the replacement fails, file remains the file's path.
            try:
                with open(file,'r') as f:
                    processed = subprocess.run(["awk", "BEGIN {ncol=0}; "
                        "($1 ~ /^[^#]/ && NF > ncol) {ncol=NF}; "
                        "{if (NF>0) {print $0} else "
                        r'{for (i=0; i<ncol; i++) printf " -"; printf "\n"}}'],
                        stdout=subprocess.PIPE, text=True,
                        input="".join(f.readlines()))
                file = io.StringIO(processed.stdout)
            except OSError:
                pass
            skip=0
            if label == 'flux': names=['i', 'psin', 'R', 'z']
        else:
            skip=HEADER_GRAY
            names=True
        try:
            run[label] = np.genfromtxt(file, skip_header=skip,
                names=names, missing_values=['-'], filling_values=np.nan)
        except OSError:
            print(f"Cannot open {file}. Skipping...")
        except ValueError:
            print(f"Bad format in {file}.")
            raise
    return run


def plot_ECprofiles(axs, profiles, label=None):
    r"""
    Plot EC absorption and ECCD profiles
    
    Parameters
    ----------
    axs : list of matplotlib Axes
        axs[0][0] : Power density
        axs[0][1] : Power
        axs[1][0] : Current density
        axs[1][1] : Current
    profiles : ndarray
        Structured numpy array containing fields with names:
        - 'rhot' : $\rho_t$, sqrt of normalized tor. flux, used as x variable
        - 'dPdV' : $dP/dV$, power density
        - 'Pins' : $P_{ins}(x) = \Int_0^x dP/dV dV/d\rho d\rho$, cumulated
            absorbed power
        - 'Jcd' : $J_{CD} = \ldots$, EC driven current density
        - 'Icdins' : $I_{CD,ins}(x) = \Int_0^x J_\phi dA/d\rho d\rho$, cumulated
            driven current

    Returns
    -------
    line_P, line_dP, line_I, line_J' : tuple of [lines2D] for the four plotted
        elements.
    """
    line_dP = axs[0][0].plot(profiles['rhot'], 1e3*profiles['dPdV'], label=label)
    line_P  = axs[0][1].plot(profiles['rhot'], profiles['Pins'], linestyle='--')
    line_J  = axs[1][0].plot(profiles['rhot'], 1e3*profiles['Jcdb'])
    line_I  = axs[1][1].plot(profiles['rhot'], 1e3*profiles['Icdins'], linestyle='--')

    return line_dP, line_P, line_J, line_I


def plot_polview(ax, rays, flux, bres,
    dpdsmax=None, dotsize=0.5, blobsize=20, marker='o', stride=1, label=None):
    r"""
    Plot rays path in 2D in a poloidal cross-section
    
    Expects to receive a structured numpy array with field names recognisable
    as spatial coordinates ('x', 'y', 'z') or ('R', 'z'). If absorption
    data are found ('alpha', 'tau'), the rays are coloured accordingly.
    
    Parameters
    ----------
    ax : matplotlib Axes
        Axes where the beam path will be plotted.
    rays : ndarray
        Structured numpy array with at least ('x', 'y', 'z') or ('R',
        'z') fields. If field names 'alpha' and 'tau' are found too, they will
        be used to shade the beam path.
    flux : ndarray
        Structured numpy array with the (R, z) contours of the flux surfaces.
    bres : ndarray
        Structured numpy array with the (R, z) contours of the EC resonances.
    dpdsmax : float, default=None
        Normalization value for the size of markers along the beam. The
        absorption markers will have size=blobsize where
        $\alpha \exp(-\tau)$ = dpdsmax.
    dotsize : float, optional, default=0.5
        Size of the dots used to represent the beam path.
    blobsize : float, optional, default=20
        Size of the markers used to represent the absorption region along the
        beam path.
    marker : optional, default='o'
        Shape of the absorption markers along the beam.
    stride : int, optional, default=1
    label : str, optional, default=None
        Unused. May be used in future for labelling.

    Returns
    -------
    dots, blobs, flux, bres : matplotlib artists representing the beam path
        with dots, the EC absorption region, the flux contours, and the EC
        resonances.
    """
    # Use colormaps with added transparency
    # Warm (yellow to red) cmap for absorption coefficient
    # Monochrome (black to white) cmap for residual power
    cmap_alpha_name = 'autumn_r'
    cmap_power_name = 'Blues'
    dP_ds_threshold = 1e-8
    try:
        map_alpha = mpl.colormaps[cmap_alpha_name + "_lin_alpha"]
    except KeyError:
        map_alpha = cmap_lin_alpha(cmap_alpha_name, register=True)
    try:
        map_power = mpl.colormaps[cmap_power_name + "_transp"]
    except KeyError:
        map_power = cmap_transp(cmap_power_name, register=True)

    if 'tau' in rays.dtype.names:
        power = np.exp(-rays['tau'])
    else:
        power = np.ones(rays['z'].shape)
    if 'alpha' in rays.dtype.names:
        dP_ds = rays['alpha']*power
    else:
        dP_ds = np.zeros(rays['z'].shape)
    
    if 'R' in rays.dtype.names:
        R = rays['R']
    elif np.all((label in rays.dtype.names for label in ('x', 'y'))):
        R = np.hypot(rays['x'],rays['y'])
    else:
        raise ValueError("At least R, or x and y, fields must be "
            "present for the poloidal cross-section plot.")

    if dpdsmax is None: dpdsmax=dP_ds.max()
    dpdsmax = max(dpdsmax, dP_ds_threshold)

    fluxlines = ax.plot(flux['R'], flux['z'], color='gray', linestyle='-', alpha=0.5)
    reslines = ax.plot(bres['R'], bres['z'], color='black', linestyle='-')
    dots = ax.scatter(R[::stride], rays['z'][::stride],
        c=power[::stride], s=dotsize,
        marker='.',
        cmap=map_power, norm=mpl.colors.Normalize(0, 1))
    blobs = ax.scatter(R, rays['z'],
        c=dP_ds, s=blobsize*(dP_ds/dpdsmax)**2, edgecolors='none',
        marker=marker,
        cmap=map_alpha, norm=mpl.colors.Normalize(0, dpdsmax))

    return dots, blobs, fluxlines, reslines


def plot_rays3d(ax, rays, dpdsmax=None, dotsize=0.5, blobsize=20, marker='o',
    stride=1, label=None):
    r"""
    Plot rays path in 3D
    
    Expects to receive a structured numpy array with field names recognisable
    as spatial coordinates ('x', 'y', 'z') or ('R', 'phi', 'z'). If absorption
    data are found ('alpha', 'tau'), the rays are coloured accordingly.
    
    Parameters
    ----------
    ax : matplotlib Axes
        Axes where the beam path will be plotted.
    rays : ndarray
        Structured numpy array with at least ('x', 'y', 'z') or ('R', 'phi',
        'z') fields. If field names 'alpha' and 'tau' are found too, they will
        be used to shade the beam path.
    dpdsmax : float, default=None
        Normalization value for the size of markers along the beam. The
        absorption markers will have size=blobsize where
        $\alpha \exp(-\tau)$ = dpdsmax.
    dotsize : float, optional, default=0.5
        Size of the dots used to represent the beam path.
    blobsize : float, optional, default=20
        Size of the markers used to represent the absorption region along the
        beam path.
    marker : optional, default='o'
        Shape of the absorption markers along the beam.
    stride : int, optional, default=1
    label : str, optional, default=None
        Unused. May be used in future for labelling.

    Returns
    -------
    dots : matplotlib artist representing the beam path with dots, shaded
        according to the residual power fraction $P/P_0 = \exp(-\tau)$.
    blobs : matplotlib artist representing the part of the beam where EC
        absorption occurs, sized and shaded according to the local rate of
        power absorption $dP/ds/P_0 = \alpha \exp(-\tau)$.
    """
    # Use colormaps with added transparency
    # Warm (yellow to red) cmap for absorption coefficient
    # Monochrome (black to white) cmap for residual power
    cmap_alpha_name = 'autumn_r'
    cmap_power_name = 'Blues'
    dP_ds_threshold = 1e-8
    try:
        map_alpha = mpl.colormaps[cmap_alpha_name + "_lin_alpha"]
    except KeyError:
        map_alpha = cmap_lin_alpha(cmap_alpha_name, register=True)
    try:
        map_power = mpl.colormaps[cmap_power_name + "_transp"]
    except KeyError:
        map_power = cmap_transp(cmap_power_name, register=True)

    if 'tau' in rays.dtype.names:
        power = np.exp(-rays['tau'])
    else:
        power = np.ones(rays['z'].shape)
    if 'alpha' in rays.dtype.names:
        dP_ds = rays['alpha']*power
    else:
        dP_ds = np.zeros(rays['z'].shape)

    if np.all((label in rays.dtype.names for label in ('x', 'y'))):
        x = rays['x']
        y = rays['y']
    elif np.all((label in rays.dtype.names for label in ('R', 'phi'))):
        x = rays['R']*np.cos(rays['phi'])
        y = rays['R']*np.sin(rays['phi'])
    else:
        raise ValueError("At least (x, y, z), or (R, phi, z), fields must be "
            "present for the 3D plot.")

    if dpdsmax is None: dpdsmax=dP_ds.max()
    dpdsmax = max(dpdsmax, dP_ds_threshold)

    dots = ax.scatter(x[::stride], y[::stride], rays['z'][::stride],
        c=power[::stride], s=dotsize,
        marker='.',
        cmap=map_power, norm=mpl.colors.Normalize(0, 1))
    blobs = ax.scatter(x, y, rays['z'],
        c=dP_ds, s=blobsize*(dP_ds/dpdsmax)**2, edgecolors='none',
        marker=marker,
        cmap=map_alpha, norm=mpl.colors.Normalize(0, dpdsmax))

    return dots, blobs


if __name__ == "__main__":
    import sys

    runs = []
    for dirname in sys.argv[1:]:
        run = read_grayout(dirname)
        if run:
            runs.append(run)
            print(f"{run['name']} loaded.")
        else:
            print(f"No valid data in {dirname}.")
    if not runs: exit()

    # Figure for profiles vs rho
    # The list axsrho has shape (2,2)
    # The axes axsrho[0][:] and axsrho[1][:] are vertically stacked and share
    # the same x axis
    # The axes axsrho[i][0] and axsrho[i][1] are overlapped with two
    # independent y axis
    figrho, axsrho = plt.subplots(2,1, sharex=True)
    axsrho[1].set_xlabel(r"$\rho_t$")
    axsrho = [[axrho, axrho.twinx()] for axrho in axsrho]
    axsrho[0][0].set_ylabel(r"$dP/dV/P_0$ [kW/m$^{-3}$/MW]")
    axsrho[0][1].set_ylabel(r"$P_{ins}/P_0$")
    axsrho[1][0].set_ylabel(r"$J_{CD}/P_0$ [kA/m$^2$/MW]")
    axsrho[1][1].set_ylabel(r"$I_{ins}/P_0$ [kA/MW]")

    plotsrho = []
    for i, run in enumerate(runs):
        plotrho = plot_ECprofiles(axsrho, run['prof'], label=run['name'])
        plotrho = {
            'name': run['name'],
            'dPdV': plotrho[0],
            'P': plotrho[1],
            'J': plotrho[2],
            'I': plotrho[3]}
        plotsrho.append(plotrho)
    figrho.legend(loc='upper center', ncols=3)

    # Figure for beam path in poloidal cross-section
    figpol, axpol = plt.subplots()
    axpol.set_xlabel(r"$R$ [m]")
    axpol.set_ylabel(r"$z$ [m]")
    axpol.set_aspect('equal')

    plotspol = []
    for i, run in enumerate(runs):
        plotpol = plot_polview(axpol, run['rays'], run['flux'], run['bres'],
            stride=5)
        plotpol = {
            'name': run['name'],
            'pow': plotpol[0],
            'abs': plotpol[1],
            'flux': plotpol[2],
            'bres': plotpol[3],
            }
        plotspol.append(plotpol)

    # Figure for 3D ray path
    fig3d, ax3d = plt.subplots(subplot_kw=dict(projection='3d'))
    xlabel3d = ax3d.set_xlabel('x [m]')
    ylabel3d = ax3d.set_ylabel('y [m]')
    zlabel3d = ax3d.set_zlabel('z [m]')
    ax3d.set_box_aspect((1, 1, 1)) # must be set BEFORE calling .set_aspect()
    limits = dict()
    for label in ('x','y','z'):
        limits[label] = (
            min([run['rays'][label].min() for run in runs]),
            max([run['rays'][label].max() for run in runs]))
#     ax.set(
#         xlim=limits['x'],
#         ylim=limits['y'],
#         zlim=limits['z'])

    dP_ds_max = max([(run['rays']['alpha']*np.exp(-run['rays']['tau'])).max()
        for run in runs])

    scatters3d = []
    for run in runs:
        scatter3d = plot_rays3d(
            ax3d, run['rays'], dP_ds_max, blobsize=2, marker='o', stride=5)
        scatter3d = {
            'name': run['name'],
            'pow': scatter3d[0],
            'abs': scatter3d[1]}
        scatters3d.append(scatter3d)
    ax3d.set_aspect('equal')

    colorbar_alpha = fig3d.colorbar(scatters3d[-1]['abs'], ax=ax3d,
        label=r"$(dP/ds)/P_0$ $[\mathrm{m}^{-1}]$", extend='max', location='left')
    colorbar_power = fig3d.colorbar(scatters3d[-1]['pow'], ax=ax3d,
        label=r"$P/P_0$", extend='neither')

    plt.show()