'''
Python interface for GRAY
'''

# output files handling
import numpy as np
import collections
import io
from configparser import ConfigParser

# standard GRAY simulation
from pathlib import Path


# Type of a parsed EQDSK file
Eqdsk = collections.namedtuple('Eqdsk', ['flux', 'boundary', 'limiter', 'q'])


def read_eqdsk(filepath: Path) -> Eqdsk:
    '''
    Loads a full G-EQDSK file and returns:
    R, z, ψ(R, z), q(ρ)
    '''
    def read(file, n):
        '''
        Reads n records of a G-EQDSK file
        '''
        x = np.empty(n)
        for i in range(n):
            x[i] = float(file.read(16).strip())
        return x

    def skip(file, n):
        '''
        Skip n records of a G-EQDSK file
        '''
        file.seek(file.tell() + 16 * n)

    with open(filepath) as file:
        file = io.StringIO(''.join(file.read().splitlines()))

    # comment string
    file.read(48)

    # number of (R, z) grid points
    _, nr, nz = [int(i) for i in file.read(4 * 3).split()]

    # build (R, z) grid
    r_dim, z_dim, _, r_left, z_mid = read(file, 5)
    r = np.linspace(r_left, r_left + r_dim, nr)
    z = np.linspace(z_mid - z_dim/2, z_mid + z_dim/2, nz)

    # ψ on the axis/boundary
    psi_min, psi_max = read(file, 5)[2:4]

    # skip other MHD quantities
    skip(file, 5 * 2)
    skip(file, nr * 4)

    # read flux and safety factor
    psi = read(file, nr * nz).reshape(nr, nz)
    q = read(file, nr)

    # normalise flux
    psi = abs(psi - psi_min) / abs(psi_max - psi_min)

    # plasma boundary and limiter contour
    nbound, nlim = [int(i) for i in file.read(5 * 2).split()]
    rbound, zbound = read(file, nbound * 2).reshape(nbound, 2).T
    rlim, zlim = read(file, nlim * 2).reshape(nlim, 2).T

    return Eqdsk(flux=(r, z, psi),
                 boundary=(rbound, zbound),
                 limiter=(rlim, zlim),
                 q=(np.sqrt(np.linspace(0, 1, nr)), abs(q)))


def read_conf(fname: Path):
    '''
    Parses gray.ini configuration file
    '''
    config = ConfigParser()
    with open(fname) as ini:
        config.read_file(ini)
    return config


def read_table(fname: Path):
    '''
    Loads a GRAY output unit file as a structured numpy array
    '''
    return np.genfromtxt(fname, skip_header=21, names=True)


def get_limiter(conf: ConfigParser, inputs: Path) -> np.array:
    '''
    Returns the limiter contour
    '''
    file = conf['equilibrium']['filenm'].strip('"')

    if conf['equilibrium'].get('iequil') == 'EQ_ANALYTICAL':
        return np.loadtxt(inputs / file, skiprows=4).T
    else:
        eqdsk = read_eqdsk(inputs / file)
        return eqdsk.limiter


def decode_index_rt(i: int) -> (str, int):
    '''
    Given the index_rt of a beam, returns the polarisation
    mode ('O' or 'X') and the number of passes
    '''
    mode = ['X', 'O'][i % 2]
    passes = int(np.floor(np.log2(1 + i)))
    return mode, passes