"""
Some utility functions used by chronometer.
"""

import os
import numpy as np
import pandas as pd
from isochrones import StarModel
from isochrones.mist import MIST_Isochrone
import h5py


def get_n_things(mods, params):
    """
    Figure out number of stars, number of global params, etc.
    """
    N, nind = len(mods), 5
    ngyro = 3
    nglob = 4
    return N, ngyro, nglob, nind


def estimate_covariance(nstars, fn):
    """
    Return the covariance matrix of the emcee samples.
    If there are more stars than the three that were used to construct this
    matrix, repeat the last five columns and rows nstar times.
    """
    with h5py.File(fn, "r") as f:
        samples = f["action_samples"][...]
    cov = np.cov(samples, rowvar=False)
    return cov


def augment(cov, N, npar):
    """
    Add the required number of parameters on to the covariance matrix.
    Repeat the individual star covariances for the last star N times.
    params:
    ------
    cov: (array)
        A 2d array of parameter covariances.
    N: (int)
        The number of stars to add on.
    npar: (int)
        The number of parameters per star.
    """
    assert N >= 0, "Number of stars should be greater than or equal to 5."
    for i in range(N):
        new_col = cov[:, -npar:]  # select last npar columns.
        aug_col = np.hstack((cov, new_col))  # Attach them to cov
        new_row = np.hstack((cov[-npar:, :], cov[-npar:, -npar:]))  # new row
        cov = np.vstack((aug_col, new_row))  # attach new row to cov.
    return cov


def mask_column(d, col, val):
    """
    Turn all the NaNs into the val.
    params:
    ------
    d: (pandas.DataFrame)
        The data
    col: (array)
        The array to mask.
    val: (float)
        The value to convert the NaNs into.
    """
    m = np.isfinite(col)
    col[~m] = np.ones(len(col[~m]))*val
    return col


def replace_nans_with_inits(d):
    """
    Turn the data into reasonable initial values.
    params:
    ------
    data: (pandas.DataFrame)
        Data read from an input file - the colours and physical parameters of
        the stars.
    If no initial mass is provided, initialise with 1 solar mass.
    If no initial feh is provided, initialise with solar feh.
    If no initial Av is provided, initialise with 0.
    """

    try:
        mass = mask_column(d, d.mass.values, 1.)
    except:
        mass = np.ones(len(d))
    d["mass"] = mass
    try:
        feh = mask_column(d, d.feh.values, 0.)
    except:
        feh = np.zeros(len(d))
    d["feh"] = feh
    try:
        Av = mask_column(d, d.Av.values, 0.)
    except:
        Av = np.zeros(len(d))
    d["Av"] = Av