import matplotlib.pyplot as plt
import pandas as pd
import os
import warnings
import sys, select
import time

def get_latest_file(file_path, suffix, v=None, warn=False):
    if v is None:
        v = 0
        name = file_path + '_' + str(v)
        while os.path.isfile(name + suffix):
            v += 1
            name = file_path + '_' + str(v)

        if os.path.isfile(file_path + '_' + str(v - 1) + suffix):
            if warn:
                warnings.warn('Found '+file_path + '_' + str(v - 1) + suffix)
            return file_path + '_' + str(v - 1) + suffix, v - 1
        else:
            raise FileNotFoundError('No file found with root name '+file_path+suffix)
    else:
        if os.path.isfile(file_path + '_'+str(v) + suffix):
            return file_path + '_' + str(v) + suffix, v
        else:
            warnings.warn('File '+file_path + '_'+str(v) + suffix+' does not exist, getting latest file with root name '+file_path)
            return get_latest_file(file_path, suffix, warn=True)

def get_empty_file_path(root_name, suffix):
    if os.path.isabs(root_name + suffix):
        if os.path.isfile(root_name + suffix):
            i = 1
            root_name = root_name + '_'+str(i)
            while os.path.isfile(root_name + suffix):
                i += 1
        return root_name + suffix
    else:
        dir_path = os.path.dirname(os.path.abspath(__file__))
        root_name = dir_path + '/' + root_name
        return get_empty_file_path(root_name, suffix)

def do_overwrite(filename, wait_time):
    if not os.path.isfile(filename+'.png'):
        return True

    print('WARNING: Overwriting', filename, 'in', wait_time, 'seconds. (Y/n)')
    print('-'*(wait_time//2))
    wait = 0
    while wait < wait_time:
        i, o, e = select.select( [sys.stdin], [], [], 2)
        if (i):
            overwrite = sys.stdin.readline().strip()
        else:
            overwrite = ''

        if overwrite == 'Y':
            return True
        elif overwrite == 'n':
            return False

        wait += 2
        print('#', end='')
    return True

def save_plot(plot_path):
    if do_overwrite(plot_path, 45):
        plt.savefig(plot_path)
        plt.clf()
        print('Saved plot to ' + plot_path)
    else:
        print('Overwrite aborted.')

def H2O_in_N2_c(v=None):
    file_path = 'output/H2O_N2_c'
    file_path, v = get_latest_file(file_path, '.csv', v=v)
    df = pd.read_csv(file_path)

    eos_list = ['VdW', 'SRK', 'PR', 'PT', 'CPA']
    for eos in eos_list:
        plt.plot(df['x_H2O'], df[eos+'_H2O'], label=eos)

    plt.legend()
    plt.xlabel(r'$x_{H2O}$')
    plt.ylabel(r'$S_T$ [K$^{-1}$]')

    plot_path = 'plots/H2O_N2_c_'+str(v)
    save_plot(plot_path)

def H2O_in_N2_T(v=None):
    file_path = 'output/H2O_N2_T'
    file_path, v = get_latest_file(file_path, '.csv', v=v)
    df = pd.read_csv(file_path)

    eos_list = ['VdW', 'SRK', 'PR', 'PT']#, 'CPA']
    for eos in eos_list:
        plt.plot(df['T [K]'], df[eos+'_H2O'], label=eos)

    plt.legend()
    plt.xlabel('T [K]')
    plt.ylabel(r'$S_T$ [K$^{-1}$]')

    plot_path = 'plots/H2O_N2_T_'+str(v)
    save_plot(plot_path)

def H2O_in_O2_c(v=None):
    file_path = 'output/H2O_O2_c'
    file_path, v = get_latest_file(file_path, '.csv', v=v)
    df = pd.read_csv(file_path)

    eos_list = ['VdW', 'SRK', 'PR', 'PT', 'CPA']
    for eos in eos_list:
        plt.plot(df['x_H2O'], df[eos+'_H2O'], label=eos)

    plt.legend()
    plt.xlabel(r'$x_{H2O}$')
    plt.ylabel(r'$S_T$ [K$^{-1}$]')

    plot_path = 'plots/H2O_O2_c_'+str(v)
    save_plot(plot_path)

def H2O_in_O2_T(v=None):
    file_path = 'output/H2O_O2_T'
    file_path, v = get_latest_file(file_path, '.csv', v=v)
    df = pd.read_csv(file_path)

    eos_list = ['VdW', 'SRK', 'PR', 'PT']#, 'CPA']
    for eos in eos_list:
        plt.plot(df['T [K]'], df[eos+'_H2O'], label=eos)

    plt.legend()
    plt.xlabel('T [K]')
    plt.ylabel(r'$S_T$ [K$^{-1}$]')

    plot_path = 'plots/H2O_O2_T_'+str(v)
    save_plot(plot_path)