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, output_lines

### 
# Script to generate BH MS systems.


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

    result_header = ['zams_mass', 'st_0', 'st_1', 'st_2', 'st_3', 'st_4', 'st_5', 'st_6', 'st_7', 'st_8', 'st_9', 'st_10', 'st_11', 'st_12', 'st_13', 'st_14', 'st_15'] 

    mass_lost_dict = {}
    for i in range(16):
        mass_lost_dict['{}'.format(i)] = 0

    total_mass_lost = 0

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

        # Check the header and act accordingly
        if (headerline=='DAVID_MASSLOSS_SN'):
            parameters = ['time', 'mass_1', 'prev_mass_1', 'zams_mass_1', 'stellar_type', 'probability']
            values = el.split()[1:]

            if not float(values[0])==0.0:
                mass_lost = float(values[2])-float(values[1])
                mass_lost_dict[values[4]] += mass_lost

                initial_mass = values[3]

                total_mass_lost += mass_lost

    result_list = [initial_mass]
    for key in mass_lost_dict.keys():
        result_list.append(str(mass_lost_dict[key]))

    result_dict = self.grid_options['result_dict']

    # This trick is necessary
    # Make the mass dict and set values
    result_dict['mass'] = result_dict.get('mass', {})
    mass_result = result_dict['mass']
    mass_result[initial_mass] = mass_result.get(initial_mass, 0) + total_mass_lost
    result_dict['mass'] = mass_result


    result_dict['probability'] = result_dict.get('probability', 0) + 0.00002123

## Set values
test_pop = Population()
test_pop.set(
    C_logging_code="""
Printf("DAVID_MASSLOSS_SN %30.12e %g %g %g %d %g\\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].stellar_type, //5
    stardata->model.probability //6
    );
""")


# Set grid variables
resolution = {'M_1': 5, 'q': 5, 'per': 5}

test_pop.add_grid_variable(
    name="lnm1",
    longname="Primary mass",
    valuerange=[1, 150],
    resolution="{}".format(resolution["M_1"]),
    spacingfunc="const(math.log(1), 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
)

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': 0.8, 'height': 1}, {'min': 0.8, 'max': 1.0, 'height': 1.0}])",
    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="logper",
    longname="log(Orbital_Period)",
    valuerange=[-2, 12],
    resolution="{}".format(resolution["per"]),
    spacingfunc="np.linspace(-2, 12, {})".format(resolution["per"]),
    precode="orbital_period = 10** logper\n", # TODO: 
    probdist="gaussian(logper,4.8, 2.3, -2.0, 12.0)",
    parameter_name="orbital_period",
    dphasevol="dln10per",
)

##########################################################################
metallicity = 0.002
test_pop.set(
    separation=1000000000, 
    orbital_period=400000000, 
    metallicity=metallicity, 
    M_1=100,
    M_2=5,

    verbose=1,
    data_dir=os.path.join(os.environ['BINARYC_DATA_ROOT'], 'testing_python', 'BHMS'),
    base_filename="BH_MS_z{}.dat".format(metallicity),
    parse_function=parse_function,
    amt_cores=2,
    )

# out = test_pop.evolve_single()
# print(out)


# quit()
test_pop.evolve_population()

def handle_output(test_pop):

    # $results is a hash reference containing
    # the results that were added up in parse_data()
    results = test_pop.grid_options['result_dict']

    print(results)

    # output the mass distribution
handle_output(test_pop)


# print(test_pop.grid_options['results_per_worker'])
# # Export settings:
# test_pop.export_all_info(use_datadir=True)

# # hdf5
# create_hdf5(test_pop.custom_options['data_dir'], name="BH_MS_z{}.hdf5".format(metallicity))