import os
import json
import time
import pickle
import sys

import matplotlib.pyplot as plt

from binarycpython.utils.grid import Population
from binarycpython.utils.functions import get_help_all, get_help, create_hdf5


import argparse


parser = argparse.ArgumentParser()
parser.add_argument(
    "amt_systems", help="the amount of systems",
)

parser.add_argument(
    "amt_nodes", help="the amount of nodes that are used for the multiprocessing",
)

parser.add_argument(
    "name_testcase", help="The name of the testcase (e.g. laptop, cluster etc)",
)


args = parser.parse_args()

amt_systems = args.amt_systems
amt_nodes = args.amt_nodes
name_testcase = args.name_testcase


## Quick script to get some output about which stars go supernova when.
def output_lines(output):
    """
    Function that outputs the lines that were recieved from the binary_c run. 
    """
    return output.splitlines()


def parse_function(self, output):
    # extract info from the population instance
    # TODO: think about whether this is smart. Passing around this object might be an overkill

    # Get some information from the
    data_dir = self.custom_options["data_dir"]
    base_filename = self.custom_options["base_filename"]

    # Check directory, make if necessary
    os.makedirs(data_dir, exist_ok=True)

    # Create filename
    outfilename = os.path.join(data_dir, base_filename)

    # Go over the output.
    for el in output_lines(output):
        headerline = el.split()[0]

        # CHeck the header and act accordingly
        if headerline == "DAVID_SN":
            parameters = ["time", "mass_1", "prev_mass_1", "zams_mass_1", "SN_type"]
            values = el.split()[1:]
            seperator = "\t"

            if not os.path.exists(outfilename):
                with open(outfilename, "w") as f:
                    f.write(seperator.join(parameters) + "\n")

            with open(outfilename, "a") as f:
                f.write(seperator.join(values) + "\n")


## Set values
test_pop = Population()
test_pop.set(
    C_logging_code="""
if(stardata->star[0].SN_type != SN_NONE)    
{
    if (stardata->model.time < stardata->model.max_evolution_time)
    {
        Printf("DAVID_SN %30.12e %g %g %g %d\\n",
            // 
            stardata->model.time, // 1
            stardata->star[0].mass, //2
            stardata->previous_stardata->star[0].mass, //3
            stardata->star[0].pms_mass, //4
            stardata->star[0].SN_type //5
      );
    };
    /* Kill the simulation to save time */
    stardata->model.max_evolution_time = stardata->model.time - stardata->model.dtm;
};
"""
)

test_pop.set(
    separation=1000000000,
    orbital_period=400000000,
    metallicity=0.002,
    data_dir=os.path.join(
        os.environ["BINARYC_DATA_ROOT"],
        "testing_python",
        "multiprocessing2",
        name_testcase,
    ),
)

res = test_pop.evolve_population_comparison(
    parse_function, amt=int(amt_systems), nodes=int(amt_nodes)
)
with open("comparison_result.dat", "a") as f:
    f.write(str(res) + "\n")


mass_distribution = np.arange(1, 200)


# evolve_population_mp(parse_function, mass_distribution)