'''
Script to quickly visualise both the inputs and output files of GRAY
Usage: python -m scripts.gray_visual -h
'''

from pathlib import Path
from matplotlib.collections import LineCollection
from matplotlib.contour import ContourSet

import argparse
import matplotlib.pyplot as plt
import numpy as np
import scripts.gray as gray


def cli_args() -> argparse.Namespace:
    '''
    Command line arguments
    '''
    parser = argparse.ArgumentParser(
        prog='gray_visual',
        description='visualises the results of a GRAY simulation')
    parser.add_argument('inputs', type=Path,
                        help='filepath of the inputs directory '
                             '(containing gray.ini)')
    parser.add_argument('outputs', type=Path,
                        help='filepath of the outputs directory')
    parser.add_argument('--kind', '-k', default=['raytracing'], nargs='+',
                        choices=['raytracing', 'beam', 'profiles', 'inputs'],
                        help='which kind of information to visualise')
    parser.add_argument('--outfile', '-o',
                        metavar='FILE', type=Path, default=None,
                        help='save the plot figure to FILE')
    parser.add_argument('--interactive', '-i', action='store_true',
                        help='show the plot in an interactive window')
    parser.add_argument('--legend',
                        action=argparse.BooleanOptionalAction, default=True,
                        help='whether to show plot legend')
    return parser.parse_args()


def align_yaxis(*axes: [plt.Axes]):
    '''
    Aligns the origins of two axes
    '''
    extrema = np.array([ax.get_ylim() for ax in axes])
    tops = extrema[:, 1] / (extrema[:, 1] - extrema[:, 0])

    # Ensure that plots (intervals) are ordered bottom to top:
    if tops[0] > tops[1]:
        axes, extrema, tops = [a[::-1] for a in (axes, extrema, tops)]

    # How much would the plot overflow if we kept current zoom levels?
    tot_span = tops[1] + 1 - tops[0]

    extrema[0, 1] = extrema[0, 0] + tot_span * (extrema[0, 1] - extrema[0, 0])
    extrema[1, 0] = extrema[1, 1] + tot_span * (extrema[1, 0] - extrema[1, 1])
    [axes[i].set_ylim(*extrema[i]) for i in range(2)]


def add_alt_axis(ax: plt.Axes, x1: np.array, x2: np.array,
                 label: str) -> plt.Axes:
    '''
    Adds an alternative x axis to the top of a plot.
    x1: existing variable
    x2: new variable
    '''
    from matplotlib.ticker import AutoMinorLocator
    from scipy.interpolate import interp1d

    ax.xaxis.set_minor_locator(AutoMinorLocator())
    ax.grid()
    f = interp1d(x1, x2, fill_value='extrapolate')
    g = interp1d(x2, x1, fill_value='extrapolate')
    alt = ax.secondary_xaxis('top', functions=(f, g))
    alt.set_xlabel(label, labelpad=8)
    alt.xaxis.set_minor_locator(AutoMinorLocator())
    return alt


def plot_circle(ax: plt.Axes, radius: float, **args) -> plt.Circle:
    '''
    Plots a circle at the origin
    '''
    c = plt.Circle((0, 0), radius, linewidth=1.5,
                   fill=False, **args)
    ax.add_patch(c)
    ax.autoscale_view()
    return c


def plot_line_color(ax: plt.Axes, x: np.array, y: np.array,
                    c: np.array, **args) -> LineCollection:
    '''
    Plots a line with variable color
    '''
    points = np.array([x, y]).T.reshape(-1, 1, 2)
    segments = np.concatenate([points[:-1], points[1:]], axis=1)

    lines = LineCollection(segments, **args)
    lines.set_array(c)
    ax.add_collection(lines)
    ax.autoscale_view()
    return lines


def plot_poloidal(inputs: Path, outputs: Path, ax: plt.Axes):
    '''
    Plots raytracing projected in the poloidal plane
    '''
    conf = gray.read_conf(inputs / 'gray.ini')
    central_ray = gray.read_table(outputs / 'central-ray.4.txt')

    ax.set_title('poloidal view', loc='right')
    ax.set_xlabel('$R$ / m')
    ax.set_ylabel('$z$ / m')

    # load flux surfaces
    surfaces = gray.read_table(outputs / 'flux-surfaces.71.txt')
    surfaces = surfaces.reshape(-1, int(surfaces['i'].max()))

    # plot plasma boundary
    bound = np.argmin(1 - surfaces['ψ_n'][:, 0])
    ax.plot(surfaces[bound]['R'], surfaces[bound]['z'],
            c='xkcd:slate grey', label='plasma boundary')

    # manual contourplot of ψ(R,z)
    for i, surf in enumerate(surfaces[:bound]):
        ax.plot(surf['R'], surf['z'], c='xkcd:grey',
                label='flux surfaces' if i == 0 else '')

    # label the rational surfaces
    labels = ['q=2', 'q=3/2', 'q=1']
    for label, surf in zip(labels, surfaces[1+bound:][::-1]):
        line = np.vstack([surf['R'], surf['z']]).T
        cont = ContourSet(ax, [1], [[line]], filled=False,
                          colors='xkcd:ocean blue',
                          linestyles='--')
        ax.clabel(cont, [1], inline=True,
                  fontsize=10, fmt={1: label})
    ax.plot(np.nan, np.nan, '--', color='xkcd:ocean blue',
            label='rational surfaces')

    # load limiter
    limiter = gray.get_limiter(conf, inputs)
    ax.plot(*limiter, c='xkcd:black', label='limiter')

    # plot cold resonance curve
    try:
        resonance = gray.read_table(outputs / 'ec-resonance.70.txt')
        for n in np.unique(resonance['n']):
            harm = resonance['n'] == n
            ax.plot(resonance['R'][harm], resonance['z'][harm],
                    c='xkcd:orange', alpha=n/resonance['n'].max(),
                    label=f"resonance $n={int(n)}$")
    except FileNotFoundError:
        pass

    # central ray
    plot_line_color(ax, central_ray['R'], central_ray['z'],
                    central_ray['α'] * np.exp(-central_ray['τ']),
                    cmap='plasma', label='rays')

    # plot outer rays
    try:
        outer_rays = gray.read_table(outputs / 'outer-rays.33.txt')
        for k in np.unique(outer_rays['k']):
            ray = outer_rays[outer_rays['k'] == k]
            plot_line_color(ax, ray['R'], ray['z'], ray['α']*np.exp(-ray['τ']),
                            cmap='plasma', alpha=0.3)
    except FileNotFoundError:
        pass


def plot_toroidal(inputs: Path, outputs: Path, ax: plt.Axes):
    '''
    Plots raytracing projected in the poloidal plane
    '''
    conf = gray.read_conf(inputs / 'gray.ini')
    central_ray = gray.read_table(outputs / 'central-ray.4.txt')

    ax.set_title('toroidal view', loc='right')
    ax.set_xlabel('$x$ / m')
    ax.set_ylabel('$y$ / m')

    limiter = gray.get_limiter(conf, inputs)

    # plot plasma boundary
    surfaces = gray.read_table(outputs / 'flux-surfaces.71.txt')
    boundary = surfaces[np.isclose(surfaces['ψ_n'], 1, 1e-3)]

    # plot plasma boundary
    plot_circle(ax, radius=boundary['R'].min(), color='xkcd:slate gray')
    plot_circle(ax, radius=boundary['R'].max(), color='xkcd:slate gray')

    # plot limiter
    if limiter[0].size > 0:
        Rmax = limiter[0].max() * 1.1
        plot_circle(ax, radius=limiter[0].min(),
                    color='xkcd:black')
        plot_circle(ax, radius=limiter[0].max(),
                    color='xkcd:black')
        # set bounds
        ax.set_xlim(-Rmax, Rmax)
        ax.set_ylim(-Rmax, Rmax)

    # plot cold resonance curve
    try:
        resonance = gray.read_table(outputs / 'ec-resonance.70.txt')
        for n in np.unique(resonance['n']):
            harm = resonance['n'] == n
            plot_circle(ax, radius=resonance['R'][harm].max(),
                        color='xkcd:orange', alpha=n/resonance['n'].max())
    except FileNotFoundError:
        pass

    # central ray
    for rt in np.unique(central_ray['index_rt']):
        ray = central_ray[central_ray['index_rt'] == rt]
        x = ray['R'] * np.cos(np.radians(ray['φ']))
        y = ray['R'] * np.sin(np.radians(ray['φ']))
        dPds = ray['α'] * np.exp(-ray['τ'])
        plot_line_color(ax, x, y, dPds, cmap='plasma')

    # outer rays
    try:
        outer_rays = gray.read_table(outputs / 'outer-rays.33.txt')
        for k in np.unique(outer_rays['k']):
            for rt in np.unique(outer_rays['index_rt']):
                ray = outer_rays[ (outer_rays['k'] == k)
                                & (outer_rays['index_rt'] == rt)]
                dPds = ray['α'] * np.exp(-ray['τ'])
                plot_line_color(ax, ray['x'], ray['y'], dPds,
                                cmap='plasma', alpha=0.3)
    except FileNotFoundError:
        pass


def plot_beam(outputs: Path, axes: [plt.Axes]):
    '''
    Plots the outermost rays in the local beam frame
    '''
    rays = gray.read_table(outputs / 'beam-shape.8.txt')
    nrays = int(max(rays['k']))

    cm = 0.01  # 1cm in m
    cmap = plt.cm.viridis

    # side view of rays, y(z)
    axes['A'].set_title('rays, side view', loc='right')
    axes['A'].set_ylabel('$y$ / cm')

    # top view of rays, x(z)
    axes['B'].set_title('rays, top view', loc='right')
    axes['B'].set_ylabel('$x$ / cm')
    axes['B'].set_xlabel('$z$ / m')

    # 3D view
    axes['C'].set_title('beam surface', loc='right')
    axes['C'].set_xlabel('$z$ / m', labelpad=11)
    axes['C'].set_ylabel('$x$ / cm', labelpad=-2)
    axes['C'].set_zlabel('$y$ / cm', labelpad=-2)
    axes['C'].set_box_aspect(aspect=(2.5, 1, 1), zoom=1.2)
    axes['C'].view_init(azim=-44)

    # plot the beam envelope
    try:
        size = gray.read_table(outputs / 'beam-size.12.txt')
        for i in ['A', 'B']:
            style = dict(c='xkcd:grey', ls=':', lw=1)
            axes[i].plot(size['s']*cm, +size['r_max'], **style)
            axes[i].plot(size['s']*cm, -size['r_max'], **style)
            axes[i].plot(size['s']*cm, +size['r_min'], **style)
            axes[i].plot(size['s']*cm, -size['r_min'], **style)
    except FileNotFoundError:
        pass

    for k in np.arange(1, nrays+1):
        ray = rays[rays['k'] == k]
        color = cmap(k / max(rays['k']))
        z = (ray['s'] + ray['z']) * cm
        # top projections
        axes['B'].plot(z, ray['x'], alpha=0.6, c=color, lw=1)
        axes['C'].plot(z, ray['x'], zs=rays['y'].min(), zdir='z',
                       alpha=0.6, c=color, lw=1)
        # side projections
        axes['A'].plot(z, ray['y'], alpha=0.8, c=color, lw=1)
        axes['C'].plot(z, ray['y'], zs=rays['x'].max(), zdir='y',
                       alpha=0.6, c=color, lw=1)

    # adjust ticks spacing
    plt.setp(axes['C'].get_yticklabels(), rotation=-25, ha='left')
    axes['C'].tick_params(axis='y', pad=-5)
    axes['C'].tick_params(axis='z', pad=0)

    # wrap arrays into 2D meshes for plot_surface
    k = rays['k'].reshape(-1, nrays)
    x = rays['x'].reshape(-1, nrays)
    y = rays['y'].reshape(-1, nrays)
    z = (rays['s'] + rays['z']).reshape(-1, nrays) * cm

    # make the xy slices closed
    x = np.column_stack([x, x[:, 0]])
    y = np.column_stack([y, y[:, 0]])
    z = np.column_stack([z, z[:, 0]])

    # plot beam surface, colored by ray indices
    axes['C'].plot_surface(z, x, y, facecolors=cmap(k/nrays), alpha=0.6)


def plot_profiles(outputs: Path, axes: [plt.Axes]):
    '''
    Plots the power and current profiles
    '''
    profiles = gray.read_table(outputs / 'ec-profiles.48.txt')

    first = profiles['index_rt'] == profiles['index_rt'].min()
    ρ_p, ρ_t = profiles[first]['ρ_p'], profiles[first]['ρ_t']

    for _, ax in axes.items():
        # set ρ_t to the x axis
        ax.set_xlabel('$ρ_t$')
        ax.set_xlim(-0.05, 1.05)
        ax.ticklabel_format(useMathText=True)

        # add secondary x axis for ρ_p
        add_alt_axis(ax, ρ_t, ρ_p, label='$ρ_p$')

    cur_dens = axes['A']
    cur_dens.set_title('current density', loc='right')
    cur_dens.set_ylabel(r'$J_\text{CD} \;/\;$ kA/m³')

    pow_dens = axes['B']
    pow_dens.set_title('power density', loc='right')
    pow_dens.set_ylabel(r'$dP/dV \;/\;$ MW/m³')

    cur_ins = axes['C']
    cur_ins.set_title('current within $ρ$', loc='right')
    cur_ins.set_ylabel(r'$I_\text{CD,inside} \;/\;$ kA')

    pow_ins = axes['D']
    pow_ins.set_title('power within $ρ$', loc='right')
    pow_ins.set_ylabel(r'$P_\text{inside} \;/\;$ MW')

    styles = [':', '-', '--', '-.']
    colors = dict(O='xkcd:orange', X='xkcd:ocean blue')

    for i in np.unique(profiles['index_rt']):
        mode, passes = gray.decode_index_rt(int(i))
        style = dict(c=colors[mode], ls=styles[passes % 4])

        slice = profiles[ (profiles['index_rt'] == i)
                        & (profiles['dPdV'] > 1e-15)]
        cur_dens.plot(slice['ρ_p'], slice['J_cd'], **style)
        pow_dens.plot(slice['ρ_p'], slice['dPdV'], **style)

        slice = profiles[ (profiles['index_rt'] == i)
                        & (profiles['P_inside'] > 1e-15)]
        cur_ins.plot(slice['ρ_p'], slice['I_cd_inside'], **style)
        pow_ins.plot(slice['ρ_p'], slice['P_inside'], **style)

    # legend
    plt.plot(np.nan, np.nan, label='O mode', c='xkcd:orange')
    plt.plot(np.nan, np.nan, label='X mode', c='xkcd:ocean blue')
    plt.plot(np.nan, np.nan, label='1° pass', c='xkcd:black', ls=styles[1])
    plt.plot(np.nan, np.nan, label='2° pass', c='xkcd:black', ls=styles[2])
    plt.plot(np.nan, np.nan, label='3° pass', c='xkcd:black', ls=styles[3])
    plt.plot(np.nan, np.nan, label='4° pass', c='xkcd:black', ls=styles[0])


def plot_inputs(inputs: Path, outputs: Path, axes: [plt.Axes]):
    '''
    Plot the input plasma profiles and MHD equilibrium
    '''
    from matplotlib.colors import TwoSlopeNorm

    maps = gray.read_table(outputs / 'inputs-maps.72.txt')

    # filter valid points
    maps = maps[maps['ψ_n'] > 0]
    flux = maps['R'], maps['z'], maps['ψ_n']

    # contour levels
    inner_levels = np.linspace(0, 0.99, 8)
    outer_levels = np.linspace(0.9991, maps['ψ_n'].max()*0.99, 6)
    all_levels = np.concatenate([inner_levels, outer_levels])

    # contour style
    norm = TwoSlopeNorm(vmin=0, vcenter=1, vmax=maps['ψ_n'].max())
    cmap = plt.cm.RdYlGn_r
    borders = dict(colors='xkcd:grey', linestyles='-', linewidths=0.8)

    # interpolated equilibrium
    interp = axes['B']
    interp.set_title('poloidal flux', loc='right')
    interp.set_xlabel('$R$ / m')
    interp.set_ylabel('$z$ / m')
    interp.tricontourf(*flux, levels=all_levels, norm=norm, cmap=cmap)
    interp.tricontour(*flux, levels=all_levels, **borders)
    interp.tricontour(*flux, levels=[1], colors='xkcd:slate gray')
    interp.plot(np.nan, np.nan, c='xkcd:slate gray', label='plasma boundary')

    # add limiter
    conf = gray.read_conf(inputs / 'gray.ini')
    limiter = gray.get_limiter(conf, inputs)
    interp.plot(*limiter, c='xkcd:black', label='limiter')

    # original MHD equilibrium
    orig = axes['A']
    if conf['equilibrium'].get('iequil') != 'EQ_ANALYTICAL':
        file = conf['equilibrium']['filenm'].strip('"')
        eqdsk = gray.read_eqdsk(inputs / file)
        orig.set_title('input G-EQDSK flux', loc='right')
        orig.set_xlabel('$R$ / m')
        orig.set_ylabel('$z$ / m')
        orig.contourf(*eqdsk.flux, levels=inner_levels, norm=norm, cmap=cmap)
        orig.contour(*eqdsk.flux, levels=inner_levels, **borders)
        orig.plot(*eqdsk.limiter, c='xkcd:black')
        orig.plot(*eqdsk.boundary, c='xkcd:slate gray')
    else:
        orig.axis('off')

    # colorbar
    bar = plt.colorbar(plt.cm.ScalarMappable(norm, cmap), cax=axes['C'])
    bar.set_label(label='normalised poloidal flux', labelpad=10)

    # Plasma radial profiles
    profiles = gray.read_table(outputs / 'kinetic-profiles.55.txt')

    axes['D'].set_title('plasma profiles', loc='right')
    add_alt_axis(axes['D'], profiles['ρ_t'], profiles['ψ_n']**0.5,
                 label='$ρ_p$')

    # density
    dens = axes['D']
    dens.set_xlabel('$ρ_t$')
    dens.set_ylabel('$n_e$ / 10²⁰m⁻³', color='xkcd:ocean blue')
    dens.plot(profiles['ρ_t'], profiles['n_e'], c='xkcd:ocean blue')

    # temperature
    temp = axes['D'].twinx()
    temp.set_ylabel('$T_e$ / keV', color='xkcd:orange')
    temp.plot(profiles['ρ_t'], profiles['T_e'], c='xkcd:orange')

    # safety factor
    inside = profiles['ρ_t'] < 1
    saft = axes['D'].twinx()
    saft.set_ylabel('$q$', color='xkcd:moss')
    saft.plot(profiles['ρ_t'][inside], profiles['q'][inside], c='xkcd:moss')
    saft.spines["right"].set_position(("axes", 1.11))

    align_yaxis(dens, temp)
    align_yaxis(temp, saft)

    # original data points
    if conf['profiles']['iprof'] == 'PROF_NUMERIC':
        # load input profiles
        orig = np.loadtxt(inputs / conf['profiles']['filenm'].strip('"'),
                          skiprows=1)

        # covert 1st column to ρ_t
        var = conf['profiles']['irho']
        if var == 'RHO_TOR':
            prof_ρ = orig[:, 0]
        elif var == 'RHO_POL':
            prof_ρ = np.interp(orig[:, 0], profiles['ψ_n']**0.5,
                               profiles['ρ_t'])
        elif var == 'RHO_PSI':
            prof_ρ = np.interp(orig[:, 0], profiles['ψ_n'], profiles['ρ_t'])

        saft_ρ = np.interp(eqdsk.q[0], profiles['ψ_n']**0.5, profiles['ρ_t'])

        # add to existing plot
        style = dict(marker='o', s=2)
        temp.scatter(prof_ρ, orig[:, 1], c='xkcd:orange', **style)
        dens.scatter(prof_ρ, orig[:, 2], c='xkcd:ocean blue', **style)
        saft.scatter(saft_ρ, eqdsk.q[1], c='xkcd:moss', **style)

        # legend
        axes['D'].plot(np.nan, np.nan, label='spline', c='k')
        axes['D'].scatter(np.nan, np.nan, label='data', c='k', **style)
        axes['D'].legend(loc='center left')


def main(kind: str, args: argparse.Namespace) -> (plt.Figure, [float]):
    '''
    Draws a single figure of a given kind
    '''
    cm = 0.3937  # 1cm in inches
    rect = [0, 0, 1, 1]  # default layout rect

    if kind == 'raytracing':
        fig, axes = plt.subplots(
            1, 2, figsize=(19.8*cm, 11*cm),
            tight_layout=True,
            subplot_kw=dict(aspect='equal'))
        plot_poloidal(args.inputs, args.outputs, axes[0])
        plot_toroidal(args.inputs, args.outputs, axes[1])

        if args.legend:
            rect = [0.2, 0, 1, 1]
            fig.tight_layout(rect=rect)
            fig.legend(loc='upper left')

    elif kind == 'beam':
        fig, axes = plt.subplot_mosaic(
            'ACD;BCD', figsize=(24*cm, 10*cm),
            per_subplot_kw={'C': dict(projection='3d')},
            width_ratios=[1, 1, 0.21])

        # manually adjust spacing (tight_layout is failing)
        axes['D'].axis('off')
        plt.subplots_adjust(hspace=0.5)
        plot_beam(args.outputs, axes)

    elif kind == 'profiles':
        fig, axes = plt.subplot_mosaic(
            'AB;CD', figsize=(23.11*cm, 13*cm),
            tight_layout=True)
        plot_profiles(args.outputs, axes)

        if args.legend:
            rect = [0.12, 0, 1, 1]
            fig.tight_layout(rect=rect)
            fig.legend(loc='upper left')

    elif kind == 'inputs':
        fig, axes = plt.subplot_mosaic(
            'ABC;DDD', figsize=(19.8*cm, 20*cm),
            tight_layout=True,
            height_ratios=[1.5, 1],
            width_ratios=[1, 1, 0.05],
            per_subplot_kw={'AB': dict(aspect='equal')})
        plot_inputs(args.inputs, args.outputs, axes)

    return fig, rect


if __name__ == '__main__':
    args = cli_args()

    def on_resize(event):
        '''
        Update layout on window resizes
        '''
        fig.tight_layout(rect=rect)
        fig.canvas.draw()

    if args.interactive:
        plt.ion()

    for kind in args.kind:
        fig, rect = main(kind, args)
        if args.interactive:
            fig.canvas.mpl_connect('resize_event', on_resize)
        if args.outfile is not None:
            if len(args.kind) > 1:
                # append plot kind to the output name
                fname = args.outfile.with_stem(args.outfile.stem + '-' + kind)
            else:
                fname = args.outfile
            fig.savefig(fname, bbox_inches='tight')

    if args.interactive:
        input('press enter to quit')