import os

# import json
# import time
# import sys

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

#########################################################
# This file serves as an example for running a population.
# The use of help(<function>) is a good way to inspect what parameters are there to use
#########################################################


## Quick script to get some output
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)

    seperator = " "

    # 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 == "MY_STELLAR_DATA":
            parameters = ["time", "mass", "zams_mass", "probability", "radius"]
            values = el.split()[1:]

            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")


# Create population object
example_pop = Population()

# If you want verbosity, set this before other things
example_pop.set(verbose=1)

# Setting values can be done via .set(<parameter_name>=<value>)
# Values that are known to be binary_c_parameters are loaded into bse_options.
# Those that are present in the default grid_options are set in grid_options
# All other values that you set are put in a custom_options dict
example_pop.set(
    # binary_c physics options
    M_1=10,  # bse_options
    separation=0,  # bse_options
    orbital_period=45000000080,  # bse_options
    max_evolution_time=15000,  # bse_options
    eccentricity=0.02,  # bse_options
    # grid_options
    amt_cores=1,  # grid_options
    verbose=1,  # verbosity. Not fully configured correctly yet but having it value of 1 prints alot of stuff
    parse_function=parse_function,  # Setting the parse function thats used in the evolve_population
    # Custom options # TODO: need to be set in grid_options probably
    data_dir=os.path.join(
        temp_dir(), "example_python_population_result"
    ),  # custom_options
    base_filename="example_pop.dat",  # custom_options
)

### Custom logging

## Below example requires changing the parse function
## very simple example of custom logging. Will work but need to change the parse function to handle that nicely.
# example_pop.set(
#     C_auto_logging={
#         "MY_HEADER_LINE": ["star[0].mass", "star[1].mass", "model.probability"]
#     }
# )


# Log the moment when the star turns into a hertzsprung-gap
example_pop.set(
    C_logging_code="""
if(stardata->star[0].stellar_type >= 2)    
{
    if (stardata->model.time < stardata->model.max_evolution_time)
    {
        Printf("MY_STELLAR_DATA %30.12e %g %g %g %g\\n",
            // 
            stardata->model.time, // 1
            stardata->star[0].mass, //2
            stardata->star[0].pms_mass, //4
            stardata->model.probability, //5
            stardata->star[0].radius // 6
      );
    };
    /* Kill the simulation to save time */
    stardata->model.max_evolution_time = stardata->model.time - stardata->model.dtm;
};
"""
)

# Add grid variables
resolution = {"M_1": 100}

# Mass
example_pop.add_grid_variable(
    name="lnm1",
    longname="Primary mass",
    valuerange=[2, 150],
    resolution="{}".format(resolution["M_1"]),
    spacingfunc="const(math.log(2), math.log(150), {})".format(resolution["M_1"]),
    precode="M_1=math.exp(lnm1)",
    probdist="three_part_powerlaw(M_1, 0.1, 0.5, 1.0, 150, -1.3, -2.3, -2.3)*M_1",
    dphasevol="dlnm1",
    parameter_name="M_1",
    condition="",  # Impose a condition on this grid variable. Mostly for a check for yourself
)

# Exporting of all the settings can be done with .export_all_info()
# on default it exports everything, but can be supressed by turning it off:
#   population settings (bse_options, grid_options, custom_options), turn off with include_population
#       settings=False
#   binary_c_defaults (all the commandline arguments that binary c accepts, and their defaults).
#       turn off with include_binary_c_defaults=False
#   include_binary_c_version_info (all the compilation info, and information about the compiled
#       parameters), turn off with include_binary_c_version_info=False
#   include_binary_c_help_all (all the help information for all the binary_c parameters),
#       turn off with include_binary_c_help_all=Fase
# On default it will write this to the custom_options['data_dir'], but that can be overriden by
#   setting use_datadir=False and providing an outfile=<>
example_pop.export_all_info()

# Creating a parsing function
# TODO: add example of setting up a parsing function

## Executing a single system
## This uses the M_1 orbital period etc set with the set function
output = example_pop.evolve_single()
print(output)

## Executing a population
## This uses the values generated by the grid_variables
# example_pop.evolve_population_mp_chunks() # TODO: update this function call

# Wrapping up the results to an hdf5 file can be done by using the create_hdf5
# (<directory containing data and settings>) This function takes the settings file
# (ending in _settings.json) and the data files (ending in .dat) from the data_dir
# and packing them into an hdf5 file, which is then written into the same data_dir directory
create_hdf5(data_dir=example_pop.custom_options["data_dir"], name="example_pop.hdf5")