import os
from binarycpython.utils.grid import Population
from binarycpython.utils.functions import (
    get_help_all,
    get_help,
    create_hdf5,
    output_lines,
)
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
#########################################################


def parse_function(self, output):
    # EXAMPLE PARSE_FUNCTION

    # extract info from the population instance

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

    parameters = ["time", "mass", "zams_mass", "probability", "radius", "stellar_type"]

    # 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":
            values = el.split()[1:]
            print(values)

            if not len(parameters) == len(values):
                print("Amount of column names isnt equal to amount of columns")
                raise ValueError

            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
    # Set companion to low mass
    M_2=0.08,  # Since in the example we run a single system, we should set the companion mass here. If we donm't do this, the code will complain.
    # grid_options
    amt_cores=2,  # grid_options
    verbose=1,  # verbosity. Not fully configured correctly yet but having it value of 1 prints alot of stuff
    # 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
)

# Creating a parsing function
example_pop.set(
    parse_function=parse_function,  # Setting the parse function thats used in the evolve_population
)

### 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 neutron
example_pop.set(
    C_logging_code="""
if(stardata->star[0].stellar_type >= 13)    
{
    if (stardata->model.time < stardata->model.max_evolution_time)
    {
        Printf("MY_STELLAR_DATA %30.12e %g %g %g %g %d\\n",
            // 
            stardata->model.time, // 1
            stardata->star[0].mass, // 2
            stardata->common.zero_age.mass[0], // 4
            stardata->model.probability, // 5
            stardata->star[0].radius, // 6
            stardata->star[0].stellar_type // 7
      );
    };
    /* Kill the simulation to save time */
    stardata->model.max_evolution_time = stardata->model.time - stardata->model.dtm;
};
"""
)

# Add grid variables
resolution = {"M_1": 20, "q": 20, "per": 40}

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

# # Mass ratio
# test_pop.add_grid_variable(
#     name="q",
#     longname="Mass ratio",
#     valuerange=["0.1/M_1", 1],
#     resolution="{}".format(resolution['q']),
#     spacingfunc="const(0.1/M_1, 1, {})".format(resolution['q']),
#     probdist="flatsections(q, [{'min': 0.1/M_1, 'max': 1.0, 'height': 1}])",
#     dphasevol="dq",
#     precode="M_2 = q * M_1",
#     parameter_name="M_2",
#     condition="",  # Impose a condition on this grid variable. Mostly for a check for yourself
# )

# #
# test_pop.add_grid_variable(
#    name="log10per", # in days
#    longname="log10(Orbital_Period)",
#    valuerange=[0.15, 5.5],
#    resolution="{}".format(resolution["per"]),
#    spacingfunc="const(0.15, 5.5, {})".format(resolution["per"]),
#    precode="""orbital_period = 10** log10per
# sep = calc_sep_from_period(M_1, M_2, orbital_period)
# sep_min = calc_sep_from_period(M_1, M_2, 10**0.15)
# sep_max = calc_sep_from_period(M_1, M_2, 10**5.5)""",
#    probdist="sana12(M_1, M_2, sep, orbital_period, sep_min, sep_max, math.log10(10**0.15), math.log10(10**5.5), -0.55)",
#    parameter_name="orbital_period",
#    dphasevol="dlog10per",
# )


# 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()

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