import numpy as np
from matplotlib import pyplot as plt
from mpl_toolkits.basemap import Basemap

def read_block(data):
    row_header = int(data[0].split(',')[1].strip())
    row_variable = int(data[0].split(',')[5].strip())
    row_column = int(data[0].split(',')[6].strip())
    row_data = int(data[0].split(',')[7].strip())

    var = data[row_header - row_variable:row_header]
    data = [tuple(s.replace('\n', '').strip() for s in row.split(',')) for row in data[row_header:-1]]

    dtype = []
    for row in var:
        row = [s.replace("'", '').strip() for s in row.split(',')]
        num = int(row[2])
        type_s = 'float' if num > 0 else 'string'
        for k in range(num):
            name_s = row[0] + ('_{:d}'.format(k + 1) if num > 1 else '')
            dtype.append((name_s, type_s))

    assert len(dtype) == row_column
    assert len(data) == row_data or row_data == -1

    data = np.array(data, dtype=dtype)

    return dtype, data

def data_read(file):
    with open(file, 'r') as f:
        data = f.readlines()

    res = []
    idx_a, idx_b = 0, 0
    while idx_b < len(data):
        if data[idx_b].replace('\n', '') == "'End of Block'":
            res.append(read_block(data[idx_a:idx_b + 1]))
            idx_a = idx_b + 1
        idx_b += 1
    res.append(read_block(data[idx_a:]))

    return res


def orbit():
    _, data = data_read('assets/orbit.csv')[0]
    lat, lon = data['Latitude'], data['Longitude'] - 180

    m = Basemap()
    m.drawcoastlines()
    m.fillcontinents(color='white', lake_color='lightskyblue')
    m.drawmapboundary(fill_color='skyblue')
    m.scatter(lon, lat, s=1)
    plt.savefig('docs/orbit.png', bbox_inches='tight')


def GCR(M=1, ax=None):
    _, data = data_read('assets/CREME86-M{:d}.csv'.format(M))[0]
    E = data['Energy']
    proton_i, proton_d = data['IFlux_1'], data['DFlux_1']
    alpha_i, alpha_d = data['IFlux_2'], data['DFlux_2']

    if ax:
        ax1 = ax
    else:
        _, ax1 = plt.subplots(1, 1)

    ax1.plot(E, proton_i, label=r'$p$')
    ax1.plot(E, alpha_i, label=r'$\alpha$')
    ax1.set_ylim([1, 5 * 1e2])
    ax1.set_yscale('log')
    ax1.set_ylabel(r'$Integrated\ Flux\ (\mathrm{m^{-2}sr^{-1}s^{-1}})$')
    ax1.legend(loc="upper left")

    ax2 = ax1.twinx()
    ax2.plot(E, proton_d, linestyle=':', label=r'$p$')
    ax2.plot(E, alpha_d, linestyle=':', label=r'$\alpha$')
    ax2.set_ylim([1e-5, 5 * 1e-2])
    ax2.set_yscale('log')
    ax2.set_ylabel(r'$Differential\ Flux\ (\mathrm{m^{-2}sr^{-1}s^{-1}MeV^{-1}})$')
    ax2.legend(loc="upper right")

    ax1.set_xlabel(r'$Energy\ (\mathrm{MeV})$')
    ax1.set_xlim(1, 1e5)
    
    if ax is None:
        plt.title('Average spectra - M = {:d}'.format(M))
        plt.xscale('log')
        plt.savefig('docs/spectra-M{:d}.png'.format(M), bbox_inches='tight')
    else:
        ax.set_title('Average spectra - M = {:d}'.format(M))
        ax.set_xscale('log')


def sun():
    _, data = data_read('assets/sun.csv')[1]
    E = data['Energy']
    proton_i, proton_d = data['IFlux_1'], data['DFlux_1']
    alpha_i, alpha_d = data['IFlux_2'], data['DFlux_2']

    _, ax1 = plt.subplots(1, 1, dpi=150)

    ax1.plot(E, proton_i, label=r'$p$')
    ax1.plot(E, alpha_i, label=r'$\alpha$')
    ax1.set_ylim([1e-2, 5 * 1e3])
    ax1.set_yscale('log')
    ax1.set_ylabel(r'$Integrated\ Flux\ (\mathrm{cm^{-2}})$')
    ax1.legend(loc="upper left")

    ax2 = ax1.twinx()
    ax2.plot(E, proton_d, linestyle=':', label=r'$p$')
    ax2.plot(E, alpha_d, linestyle=':', label=r'$\alpha$')
    ax2.set_ylim([1e-5, 5 * 1e0])
    ax2.set_yscale('log')
    ax2.set_ylabel(r'$Differential\ Flux\ (\mathrm{cm^{-2}MeV^{-1}})$')
    ax2.legend(loc="upper right")

    ax1.set_xlabel(r'$Energy\ (\mathrm{MeV})$')
    ax1.set_xlim(1e-1, 1e3)

    plt.title('Average spectra - Sun')
    plt.xscale('log')
    plt.savefig('docs/spectra-sun', bbox_inches='tight')


# orbit()

# fig = plt.figure(figsize=(20, 10), dpi=150)

# GCR(1, fig.add_subplot(1, 2, 1))
# GCR(3, fig.add_subplot(1, 2, 2))

# fig.savefig('docs/spectra.png', bbox_inches='tight')

# GCR(1)
# GCR(3)

sun()