#!/usr/bin/env python
"""
Plot the models in the Zenodo of SC-REM paper.

:copyright:
    Johannes Kemper (johannes.kemper@erdw.ethz.ch), 2023
:license:
    GNU Lesser General Public License, Version 3
    (http://www.gnu.org/copyleft/lgpl.html)
"""
import matplotlib.pyplot as plt
import numpy as np


if __name__ == '__main__':
    # load cis
    cidata = {}
    cidata[0.25] = [np.genfromtxt('screm-25p-low.dat',
                                  skip_header=1,
                                  delimiter=',',
                                  unpack=True),
                    np.genfromtxt('screm-25p-high.dat',
                                  skip_header=1,
                                  delimiter=',',
                                  unpack=True)]
    cidata[0.5] = [np.genfromtxt('screm-50p-low.dat',
                                 skip_header=1,
                                 delimiter=',',
                                 unpack=True),
                   np.genfromtxt('screm-50p-high.dat',
                                 skip_header=1,
                                 delimiter=',',
                                 unpack=True)]
    cidata[0.75] = [np.genfromtxt('screm-75p-low.dat',
                                  skip_header=1,
                                  delimiter=',',
                                  unpack=True),
                    np.genfromtxt('screm-75p-high.dat',
                                  skip_header=1,
                                  delimiter=',',
                                  unpack=True)]

    # set plot labels and other preliminaries
    xlabels = [r'$\rho$ [g/cm**3]', 'VP [km/s]', 'VS [km/s]',
               r'$Q_\kappa$', r'$Q_\mu$', 'K [GPa]', r'$\mu$ [GPa]',
               'pressure [GPa]', 'temperature [K]']
    fluid_indices = np.where(cidata[0.75][0][4] < 1e-5)[0]

    # set plot_indiv to plot individual parameters
    # select from dictionary par_dict
    user_input = input('Do you want to plot individual parameters (y/n): ')
    if user_input.lower() in ['y', 'yes', 'ye']:
        plot_indiv = True
    elif user_input.lower() in ['n', 'no']:
        plot_indiv = False
    else:
        print('Input either y or n.')
        exit()

    if plot_indiv:
        par_dict = {'rho': 2,
                    'vp': 3, 'vs': 4,
                    'Qkappa': 5, 'Qmu': 6,
                    'K': 7, 'mu': 8,
                    'p': 9, 'T': 10}

        print(f'Possible parameter options are: {list(par_dict.keys())}')
        ptype = input('Which parameter do you want to plot: ')
        ptype = ptype.replace('\'', '')

        try:
            e = par_dict[ptype]
        except Exception:
            print('Entered from parameter type in variable ptype.')
            print(f'Possible options are: {list(par_dict.keys())}')
            exit()

        fluid_indices = np.where(cidata[0.75][0][par_dict['vs']] < 1e-5)[0]

        xlog = [3, 4]
        fig, ax = plt.subplots(1, 1)
        ax.set_ylabel('radius [km]')
        ax.set_ylim([0.0, 6371.0])

        cfs = [0.75, 0.5, 0.25]
        ax.set_xlabel(xlabels[e - 2])
        if (e - 2) in xlog:
            ax.set_xscale('log')
        for i, cf in enumerate(cfs):
            ax.fill_betweenx(cidata[cf][0][0],
                             cidata[cf][0][e],
                             cidata[cf][1][e],
                             color='tab:blue',
                             alpha=1.0 - cf)
        ax.legend()
    else:
        xlog = [3, 4]
        fig, axs = plt.subplots(1, len(xlabels), sharey=True)
        axs[0].set_ylabel('radius [km]')
        axs[0].set_ylim([0.0, 6371.0])

        cfs = [0.75, 0.5, 0.25]
        for e, ax in enumerate(axs, start=2):
            ax.set_xlabel(xlabels[e - 2])
            if (e - 2) in xlog:
                ax.set_xscale('log')
            for cf in cfs:
                ax.fill_betweenx(cidata[cf][0][0],
                                 cidata[cf][0][e],
                                 cidata[cf][1][e],
                                 color='tab:blue',
                                 alpha=1.0 - cf)
        axs[-1].legend()

    plt.tight_layout()
    plt.show()