import numpy as np
import warnings
import pandas as pd

from ... import utils

from .sampler import register_sampler
from .. import InitialBinaryTable
from ... import _evolvebin

__author__ = "Katelyn Breivik <katie.breivik@gmail.com>"
__credits__ = ("Scott Coughlin <scott.coughlin@ligo.org>, Michael Zevin <michael.j.zevin@gmail.com>, "
               "Tom Wagg <tomjwagg@gmail.com>")
__all__ = ["get_independent_sampler", "Sample"]

def get_independent_sampler(
    final_kstar1,
    final_kstar2,
    primary_model,
    ecc_model,
    porb_model,
    SF_start,
    SF_duration,
    binfrac_model,
    met,
    size=None,
    total_mass=np.inf,
    sampling_target="size",
    trim_extra_samples=False,
    **kwargs
):
    """Generates an initial binary sample according to user specified models

    Parameters
    ----------
    final_kstar1 : `int or list`
        Int or list of final kstar1

    final_kstar2 : `int or list`
        Int or list of final kstar2

    primary_model : `str`
        Model to sample primary mass; choices include: kroupa93, kroupa01, salpeter55, custom
        if 'custom' is selected, must also pass arguemts:
        alphas : `array`
            list of power law indicies
        mcuts : `array`
            breaks in the power laws.
        e.g. alphas=[-1.3,-2.3,-2.3],mcuts=[0.08,0.5,1.0,150.] reproduces standard Kroupa2001 IMF

    ecc_model : `str`
        Model to sample eccentricity; choices include: thermal, uniform, sana12

    porb_model : `str`
        Model to sample orbital period; choices include: log_uniform, sana12, raghavan10, moe19

    qmin : `float`
        kwarg which sets the minimum mass ratio for sampling the secondary
        where the mass ratio distribution is flat in q
        if q > 0, qmin sets the minimum mass ratio
        q = -1, this limits the minimum mass ratio to be set such that
        the pre-MS lifetime of the secondary is not longer than the full
        lifetime of the primary if it were to evolve as a single star

    m_max : `float`
        kwarg which sets the maximum primary and secondary mass for sampling
        NOTE: this value changes the range of the IMF and should *not* be used
            as a means of selecting certain kstar types!

    m1_min : `float`
        kwarg which sets the minimum primary mass for sampling
        NOTE: this value changes the range of the IMF and should *not* be used
            as a means of selecting certain kstar types!

    m2_min : `float`
        kwarg which sets the minimum secondary mass for sampling
        the secondary as uniform in mass_2 between m2_min and mass_1

    msort : `float`
        Stars with M>msort can have different pairing and sampling of companions

    qmin_msort : `float`
        Same as qmin for M>msort

    m2_min_msort : `float`
        Same as m2_min for M>msort

    SF_start : `float`
        Time in the past when star formation initiates in Myr

    SF_duration : `float`
        Duration of constant star formation beginning from SF_Start in Myr

    binfrac_model : `str or float`
        Model for binary fraction; choices include: vanHaaften or a fraction where 1.0 is 100% binaries

    binfrac_model_msort : `str or float`
        Same as binfrac_model for M>msort

    met : `float`
        Sets the metallicity of the binary population where solar metallicity is zsun

    size : `int`
        Size of the population to sample

    total_mass : `float`
        Total mass to use as a target for sampling

    sampling_target : `str`
        Which type of target to use for sampling (either "size" or "total_mass"), by default "size".
        Note that `total_mass` must not be None when `sampling_target=="total_mass"`.

    trim_extra_samples : `str`
        Whether to trim the sampled population so that the total mass sampled is as close as possible to
        `total_mass`. Ignored when `sampling_target==size`.
        Note that given the discrete mass of stars, this could mean your sample is off by 300
        solar masses in the worst case scenario (of a 150+150 binary being sampled). In reality the majority
        of cases track the target total mass to within a solar mass.

    zsun : `float`
        optional kwarg for setting effective radii, default is 0.02

    Returns
    -------
    InitialBinaryTable : `pandas.DataFrame`
        DataFrame in the format of the InitialBinaryTable

    mass_singles : `float`
        Total mass in single stars needed to generate population

    mass_binaries : `float`
        Total mass in binaries needed to generate population

    n_singles : `int`
        Number of single stars needed to generate a population

    n_binaries : `int`
        Number of binaries needed to generate a population
    """
    if sampling_target == "total_mass" and (total_mass is None or total_mass == np.inf):
        raise ValueError("If `sampling_target == 'total mass'` then `total_mass` must be supplied")
    if size is None and (total_mass is None or total_mass == np.inf):
        raise ValueError("Either a sample `size` or `total_mass` must be supplied")
    elif size is None:
        size = int(total_mass)

    if binfrac_model == 0.0 and sampling_target == "size":
        raise ValueError(("If `binfrac_model == 0.0` then `sampling_target` must be 'total_mass'. Otherwise "
                          "you are targetting a population of `size` binaries but will never select any."))

    final_kstar1 = [final_kstar1] if isinstance(final_kstar1, (int, float)) else final_kstar1
    final_kstar2 = [final_kstar2] if isinstance(final_kstar2, (int, float)) else final_kstar2
    primary_min, primary_max, secondary_min, secondary_max = utils.mass_min_max_select(
        final_kstar1, final_kstar2, **kwargs)
    initconditions = Sample()

    # set up multiplier if the mass sampling is inefficient
    multiplier = 1

    # track samples to actually return (after masks)
    mass1_singles = []
    mass1_binary = []
    mass2_binary = []
    binfrac = []

    # track the total mass of singles and binaries sampled
    m_sampled_singles = 0.0
    m_sampled_binaries = 0.0

    # track the total number of stars sampled
    n_singles = 0
    n_binaries = 0

    # if porb_model = `moe19`, the binary fraction is fixed based on the metallicity
    if porb_model == "moe19":
        binfrac_model = utils.get_met_dep_binfrac(met)
        warnings.warn('your supplied binfrac_model has been overwritten to {} match Moe+2019'.format(binfrac_model))

    # define a function that evaluates whether you've reached your sampling target
    target = lambda mass1_binary, size, m_sampled_singles, m_sampled_binaries, total_mass:\
        len(mass1_binary) < size if sampling_target == "size" else m_sampled_singles + m_sampled_binaries < total_mass

    # sample until you've reached your target
    while target(mass1_binary, size, m_sampled_singles, m_sampled_binaries, total_mass):
        # sample primary masses
        mass1, _ = initconditions.sample_primary(primary_model, size=int(size * multiplier), **kwargs)

        # split them into binaries or single stars
        (mass1_binaries, mass_single, binfrac_binaries, binary_index,
        ) = initconditions.binary_select(mass1, binfrac_model=binfrac_model, **kwargs)

        # sample secondary masses for the single stars
        mass2_binaries = initconditions.sample_secondary(mass1_binaries, **kwargs)

        # check if this batch of samples will take us over our sampling target
        if not target(mass1_binary, size,
                      m_sampled_singles + np.sum(mass_single),
                      m_sampled_binaries + np.sum(mass1_binaries) + np.sum(mass2_binaries),
                      total_mass) and trim_extra_samples and sampling_target == "total_mass":
            # get the cumulative total mass of the samples
            total_mass_list = np.copy(mass1)
            total_mass_list[binary_index] += mass2_binaries
            sampled_so_far = m_sampled_singles + m_sampled_binaries
            cumulative_total_mass = sampled_so_far + np.cumsum(total_mass_list)

            # find the boundary for reaching the right total mass
            threshold_index = np.where(cumulative_total_mass > total_mass)[0][0]

            keep_offset = abs(cumulative_total_mass[threshold_index] - total_mass)
            drop_offset = abs(cumulative_total_mass[threshold_index - 1] - total_mass)
            lim = threshold_index - 1 if (keep_offset > drop_offset) else threshold_index

            # work out how many singles vs. binaries to delete
            one_if_binary = np.zeros(len(mass1))
            one_if_binary[binary_index] = 1
            sb_delete = one_if_binary[lim + 1:]
            n_single_delete = (sb_delete == 0).sum()
            n_binary_delete = (sb_delete == 1).sum()

            # delete em!
            if n_single_delete > 0:
                mass_single = mass_single[:-n_single_delete]
            if n_binary_delete > 0:
                mass1_binaries = mass1_binaries[:-n_binary_delete]
                mass2_binaries = mass2_binaries[:-n_binary_delete]
                binfrac_binaries = binfrac_binaries[:-n_binary_delete]

            # ensure we don't loop again after this
            target = lambda mass1_binary, size, m_sampled_singles, m_sampled_binaries, total_mass: False

        # track the mass sampled
        m_sampled_singles += sum(mass_single)
        m_sampled_binaries += sum(mass1_binaries)
        m_sampled_binaries += sum(mass2_binaries)

        # track the total number sampled
        n_singles += len(mass_single)
        n_binaries += len(mass1_binaries)

        # select out the primaries and secondaries that will produce the final kstars
        ind_select = (  (mass1_binaries > primary_min)
                      & (mass1_binaries < primary_max)
                      & (mass2_binaries > secondary_min)
                      & (mass2_binaries < secondary_max))
        mass1_binary.extend(mass1_binaries[ind_select])
        mass2_binary.extend(mass2_binaries[ind_select])
        binfrac.extend(binfrac_binaries[ind_select])

        # select out the single stars that will produce the final kstar
        mass1_singles.extend(mass_single[(mass_single > primary_min) & (mass_single < primary_max)])

        # check to see if we should increase the multiplier factor to sample the population more quickly
        if target(mass1_binary, size / 100, m_sampled_singles, m_sampled_binaries, total_mass / 100):
            # well this sampling rate is clearly not working time to increase
            # the multiplier by an order of magnitude
            multiplier *= 10

    mass1_binary = np.array(mass1_binary)
    mass2_binary = np.array(mass2_binary)
    binfrac = np.asarray(binfrac)
    mass1_singles = np.asarray(mass1_singles)

    zsun = kwargs.pop("zsun", 0.02)

    rad1 = initconditions.set_reff(mass1_binary, metallicity=met, zsun=zsun)
    rad2 = initconditions.set_reff(mass2_binary, metallicity=met, zsun=zsun)

    # sample periods and eccentricities
    # if the porb_model is moe19, the metallicity needs to be supplied
    if porb_model == "moe19":
        porb,aRL_over_a = initconditions.sample_porb(
            mass1_binary, mass2_binary, rad1, rad2, porb_model, met=met, size=mass1_binary.size
        )
    else:
        porb,aRL_over_a = initconditions.sample_porb(
            mass1_binary, mass2_binary, rad1, rad2, porb_model, size=mass1_binary.size
        )
    ecc = initconditions.sample_ecc(aRL_over_a, ecc_model, size=mass1_binary.size)

    tphysf, metallicity = initconditions.sample_SFH(
        SF_start=SF_start, SF_duration=SF_duration, met=met, size=mass1_binary.size
    )
    metallicity[metallicity < 1e-4] = 1e-4
    metallicity[metallicity > 0.03] = 0.03
    kstar1 = initconditions.set_kstar(mass1_binary)
    kstar2 = initconditions.set_kstar(mass2_binary)

    if kwargs.pop("keep_singles", False):
        binary_table = InitialBinaryTable.InitialBinaries(
            mass1_binary,
            mass2_binary,
            porb,
            ecc,
            tphysf,
            kstar1,
            kstar2,
            metallicity,
            binfrac=binfrac,
        )
        tphysf_singles, metallicity_singles = initconditions.sample_SFH(
            SF_start=SF_start, SF_duration=SF_duration, met=met, size=mass1_singles.size
        )
        metallicity_singles[metallicity_singles < 1e-4] = 1e-4
        metallicity_singles[metallicity_singles > 0.03] = 0.03
        kstar1_singles = initconditions.set_kstar(mass1_singles)
        singles_table = InitialBinaryTable.InitialBinaries(
            mass1_singles,                          # mass1
            np.ones_like(mass1_singles) * 0,        # mass2 (all massless remnants)
            np.ones_like(mass1_singles) * -1,       # porb (single not binary)
            np.ones_like(mass1_singles) * -1,       # ecc (single not binary)
            tphysf_singles,                         # tphysf
            kstar1_singles,                         # kstar1
            np.ones_like(mass1_singles) * 15,       # kstar2 (all massless remnants)
            metallicity_singles,                    # metallicity
        )
        binary_table = pd.concat([binary_table, singles_table])
    else:
        binary_table = InitialBinaryTable.InitialBinaries(
            mass1_binary,
            mass2_binary,
            porb,
            ecc,
            tphysf,
            kstar1,
            kstar2,
            metallicity,
            binfrac=binfrac,
        )

    return (
        binary_table,
        m_sampled_singles,
        m_sampled_binaries,
        n_singles,
        n_binaries,
    )

register_sampler(
    "independent",
    InitialBinaryTable,
    get_independent_sampler,
    usage="final_kstar1, final_kstar2, binfrac_model, primary_model, ecc_model, SFH_model, component_age, metallicity, size",
)

class Sample(object):
    # sample primary masses
    def sample_primary(self, primary_model='kroupa01', size=None, **kwargs):
        """Sample the primary mass (always the most massive star) from a user-selected model

            kroupa93 follows Kroupa (1993), normalization comes from
            `Hurley 2002 <https://arxiv.org/abs/astro-ph/0201220>`_
            between 0.08 and 150 Msun
            salpter55 follows
            `Salpeter (1955) <http://adsabs.harvard.edu/abs/1955ApJ...121..161S>`_
            between 0.08 and 150 Msun
            kroupa01 follows Kroupa (2001) <https://arxiv.org/abs/astro-ph/0009005>
            between 0.08 and 100 Msun

            Parameters
            ----------
            primary_model : str, optional
            model for mass distribution; choose from:

            kroupa93 follows Kroupa (1993), normalization comes from
            `Hurley 2002 <https://arxiv.org/abs/astro-ph/0201220>`_
            valid for masses between 0.1 and 100 Msun

            salpter55 follows
            `Salpeter (1955) <http://adsabs.harvard.edu/abs/1955ApJ...121..161S>`_
            valid for masses between 0.1 and 100 Msun

            kroupa01 follows Kroupa (2001), normalization comes from
            `Hurley 2002 <https://arxiv.org/abs/astro-ph/0009005>`_
            valid for masses between 0.1 and 100 Msun

            custom is a generic piecewise power law that takes in the power
            law slopes and break points given in the optional input lists (alphas, mcuts)
            default alphas and mcuts yield an IMF identical to kroupa01

            Default kroupa01

            size : int, optional
            number of initial primary masses to sample
            NOTE: this is set in cosmic-pop call as Nstep

            alphas : array, optional
            absolute values of the power law slopes for primary_model = 'custom'
            Default [-1.3,-2.3,-2.3] (identical to slopes for primary_model = 'kroupa01')

            mcuts : array, optional, units of Msun
            break points separating the power law 'pieces' for primary_model = 'custom'
            Default [0.08,0.5,1.0,150.] (identical to breaks for primary_model = 'kroupa01')

            Optional kwargs are defined in `get_independent_sampler`

            Returns
            -------
            a_0 : array
            Sampled primary masses
            np.sum(a_0) : float
            Total amount of mass sampled
            """

        # Read in m1_min and m_max kwargs, if provided
        m1_min = kwargs["m1_min"] if "m1_min" in kwargs.keys() else 0.08
        m_max = kwargs["m_max"] if "m_max" in kwargs.keys() else 150.0

        # Make sure m1_min value is below 0.5, since otherwise it will not work for Kroupa IMF
        if m1_min > 0.5:
            raise ValueError("m1_min must be greater than 0.5 Msun")

        if primary_model == 'kroupa93':
            alphas, mcuts = [-1.3,-2.2,-2.7], [m1_min,0.5,1.0,m_max]
        # Since COSMIC/BSE can't handle < 0.08Msun, we by default truncate at 0.08 Msun instead of 0.01
        elif primary_model == 'kroupa01':
            alphas, mcuts = [-1.3,-2.3], [m1_min,0.5,m_max]
        elif primary_model == 'salpeter55':
            alphas, mcuts = [-2.35], [m1_min,m_max]
        elif primary_model == 'custom':
            if 'alphas' in kwargs and 'mcuts' in kwargs:
                alphas = kwargs.pop("alphas", [-1.3,-2.3,-2.3])
                mcuts = kwargs.pop("mcuts", [m1_min,0.5,1.0,m_max])
            else:
                raise ValueError("You must supply both alphas and mcuts to use"
                                 " a custom IMF generator")

        Ncumulative, Ntotal, coeff = [], 0., 1.
        for i in range(len(alphas)):
            g = 1. + alphas[i]
            # Compute this piece of the IMF's contribution to Ntotal
            if alphas[i] == -1: Ntotal += coeff * np.log(mcuts[i+1]/mcuts[i])
            else: Ntotal += coeff/g * (mcuts[i+1]**g - mcuts[i]**g)
            Ncumulative.append(Ntotal)
            if i < len(alphas)-1: coeff *= mcuts[i+1]**(-alphas[i+1]+alphas[i])

        cutoffs = np.array(Ncumulative)/Ntotal
        u = np.random.uniform(0.,1.,size)
        idxs = [() for i in range(len(alphas))]

        for i in range(len(alphas)):
            if i == 0: idxs[i], = np.where(u <= cutoffs[0])
            elif i < len(alphas)-1: idxs[i], = np.where((u > cutoffs[i-1]) & (u <= cutoffs[i]))
            else: idxs[i], = np.where(u > cutoffs[i-1])
        for i in range(len(alphas)):
            if alphas[i] == -1.0:
                u[idxs[i]] = 10**np.random.uniform(np.log10(mcuts[i]),
                                                   np.log10(mcuts[i+1]),
                                                   len(idxs[i]))
            else:
                u[idxs[i]] = utils.rndm(a=mcuts[i], b=mcuts[i+1], g=alphas[i], size=len(idxs[i]))

        return u, np.sum(u)

    # sample secondary mass
    def sample_secondary(self, primary_mass, **kwargs):
        """Sample a secondary mass using draws from a uniform mass ratio distribution motivated by
        `Mazeh et al. (1992) <http://adsabs.harvard.edu/abs/1992ApJ...401..265M>`_
        and `Goldberg & Mazeh (1994) <http://adsabs.harvard.edu/abs/1994ApJ...429..362G>`_

        NOTE: the lower lim is set by either qmin or m2_min which are passed as kwargs

        Parameters
        ----------
        primary_mass : `array`
            sets the maximum secondary mass (for a maximum mass ratio of 1)

        Optional kwargs are defined in `get_independent_sampler`

        Returns
        -------
        secondary_mass : array
            sampled secondary masses with array size matching size of
            primary_mass
        """

        qmin = kwargs["qmin"] if "qmin" in kwargs.keys() else 0.0
        m1_min = kwargs["m1_min"] if "m1_min" in kwargs.keys() else 0.08
        m2_min = kwargs["m2_min"] if "m2_min" in kwargs.keys() else None
        if (m2_min is None) & (qmin is None):
            warnings.warn("It is highly recommended that you specify either qmin or m2_min!")
        if (m2_min is not None) and (m2_min > m1_min):
            raise ValueError("The m2_min you specified is above the minimum"
                             " primary mass of the IMF, either lower m2_min or"
                             " raise the lower value of your sampled primaries")

        # --- `msort` kwarg can be set to have different qmin above `msort`
        msort = kwargs["msort"] if "msort" in kwargs.keys() else None
        qmin_msort = kwargs["qmin_msort"] if "qmin_msort" in kwargs.keys() else None
        m2_min_msort = kwargs["m2_min_msort"] if "m2_min_msort" in kwargs.keys() else None
        if (msort is None) and (qmin_msort is not None):
            raise ValueError("If qmin_msort is specified, you must also supply a value for msort")
        if (msort is None) and (m2_min_msort is not None):
            raise ValueError("If m2_min_msort is specified, you must also supply a value for msort")
        if (m2_min_msort is not None) and (m2_min_msort > msort):
            raise ValueError("The m2_min_msort you specified is above the minimum"
                             " primary mass of the high-mass binaries msort")

        if (msort is not None) and (qmin_msort is not None):
            (highmassIdx,) = np.where(primary_mass >= msort)
            (lowmassIdx,) = np.where(primary_mass < msort)
        else:
            (highmassIdx,) = np.where(primary_mass < 0)
            (lowmassIdx,) = np.where(primary_mass >= 0)   # all idxs

        qmin_vals = -1 * np.ones_like(primary_mass)
        # --- qmin for low-mass systems (all systems if msort is not specified)
        if (qmin > 0.0):
            qmin_vals[lowmassIdx] = qmin * np.ones_like(primary_mass[lowmassIdx])
        elif (qmin < 0.0):
            # mass-dependent qmin, assume qmin=0.1 for m_primary<5
            dat = np.array([[5.0, 0.1363522012578616],
                            [6.999999999999993, 0.1363522012578616],
                            [12.599999999999994, 0.11874213836477984],
                            [20.999999999999993, 0.09962264150943395],
                            [29.39999999999999, 0.0820125786163522],
                            [41, 0.06490566037735851],
                            [55, 0.052327044025157254],
                            [70.19999999999999, 0.04301886792452836],
                            [87.4, 0.03622641509433966],
                            [107.40000000000002, 0.030188679245283068],
                            [133.40000000000003, 0.02515723270440262],
                            [156.60000000000002, 0.02163522012578628],
                            [175.40000000000003, 0.01962264150943399],
                            [200.20000000000005, 0.017358490566037776]])
            from scipy.interpolate import interp1d
            qmin_interp = interp1d(dat[:, 0], dat[:, 1])
            qmin_vals[lowmassIdx] = np.ones_like(primary_mass[lowmassIdx]) * 0.1
            ind_5, = np.where(primary_mass[lowmassIdx] > 5.0)
            qmin_vals[lowmassIdx][ind_5] = qmin_interp(primary_mass[lowmassIdx][ind_5])
        else:
            qmin_vals[lowmassIdx] = np.zeros_like(primary_mass[lowmassIdx])
        # --- qmin for high-mass systems, if msort and qmin_msort are specified
        if (msort is not None) and (qmin_msort is not None):
            if (qmin_msort > 0.0):
                qmin_vals[highmassIdx] = qmin_msort * np.ones_like(primary_mass[highmassIdx])
            elif (qmin_msort < 0.0):
                # mass-dependent qmin, assume qmin=0.1 for m_primary<5
                dat = np.array([[5.0, 0.1363522012578616],
                                [6.999999999999993, 0.1363522012578616],
                                [12.599999999999994, 0.11874213836477984],
                                [20.999999999999993, 0.09962264150943395],
                                [29.39999999999999, 0.0820125786163522],
                                [41, 0.06490566037735851],
                                [55, 0.052327044025157254],
                                [70.19999999999999, 0.04301886792452836],
                                [87.4, 0.03622641509433966],
                                [107.40000000000002, 0.030188679245283068],
                                [133.40000000000003, 0.02515723270440262],
                                [156.60000000000002, 0.02163522012578628],
                                [175.40000000000003, 0.01962264150943399],
                                [200.20000000000005, 0.017358490566037776]])
                from scipy.interpolate import interp1d
                qmin_interp = interp1d(dat[:, 0], dat[:, 1])
                qmin_vals[highmassIdx] = np.ones_like(primary_mass[highmassIdx]) * 0.1
                ind_5, = np.where(primary_mass[highmassIdx] > 5.0)
                qmin_vals[highmassIdx][ind_5] = qmin_interp(primary_mass[highmassIdx][ind_5])
            else:
                qmin_vals[highmassIdx] = np.zeros_like(primary_mass[highmassIdx])

        # --- apply m2_min and m2_min_msort, if specified
        if m2_min is not None:
            qmin_vals[lowmassIdx] = np.maximum(qmin_vals[lowmassIdx], m2_min/primary_mass[lowmassIdx])
        if m2_min_msort is not None:
            qmin_vals[highmassIdx] = np.maximum(qmin_vals[highmassIdx], m2_min_msort/primary_mass[highmassIdx])

        # --- now, randomly sample mass ratios and get secondary masses
        secondary_mass = np.random.uniform(qmin_vals, 1) * primary_mass

        return secondary_mass

    def binary_select(self, primary_mass, binfrac_model=0.5, **kwargs):
        """Select which primary masses will have a companion using
        either a binary fraction specified by a float or a
        primary-mass dependent binary fraction following
        `van Haaften et al.(2009) <http://adsabs.harvard.edu/abs/2013A%26A...552A..69V>`_ in appdx

        Parameters
        ----------
            primary_mass : array
                Mass that determines the binary fraction

            binfrac_model : str or float
                vanHaaften - primary mass dependent and ONLY VALID up to 100 Msun
                float - fraction of binaries; 0.5 means 2 in 3 stars are a binary pair while 1
                means every star is in a binary pair

        Optional kwargs are defined in `get_independent_sampler`

        Returns
        -------
            stars_in_binary : array
                primary masses that will have a binary companion

            stars_in_single : array
                primary masses that will be single stars

            binary_fraction : array
                system-specific probability of being in a binary

            binaryIdx : array
                Idx of stars in binary
        """

        # --- `msort` kwarg can be set to have different binary fraction above `msort`
        msort = kwargs["msort"] if "msort" in kwargs.keys() else None
        binfrac_model_msort = kwargs["binfrac_model_msort"] if "binfrac_model_msort" in kwargs.keys() else None
        if (msort is None) and (binfrac_model_msort is not None):
            raise ValueError("If binfrac_model_msort is specified, you must also supply a value for msort")
        if (msort is not None) and (binfrac_model_msort is not None):
            (highmassIdx,) = np.where(primary_mass >= msort)
            (lowmassIdx,) = np.where(primary_mass < msort)
        else:
            (highmassIdx,) = np.where(primary_mass < 0)
            (lowmassIdx,) = np.where(primary_mass >= 0)   # all idxs

        # --- read in binfrac models
        if type(binfrac_model) == str:
            if binfrac_model == "vanHaaften":
                binary_fraction_low = 1 / 2.0 + 1 / \
                    4.0 * np.log10(primary_mass[lowmassIdx])
                binary_choose_low = np.random.uniform(
                    0, 1.0, primary_mass[lowmassIdx].size)

                (singleIdx_low,) = np.where(
                    binary_fraction_low < binary_choose_low)
                (binaryIdx_low,) = np.where(
                    binary_fraction_low >= binary_choose_low)
            else:
                raise ValueError(
                    "You have supplied a non-supported binary fraction model. Please choose vanHaaften or a float"
                )
        elif type(binfrac_model) == float:
            if (binfrac_model <= 1.0) & (binfrac_model >= 0.0):
                binary_fraction_low = binfrac_model * \
                    np.ones(primary_mass[lowmassIdx].size)
                binary_choose_low = np.random.uniform(
                    0, 1.0, primary_mass[lowmassIdx].size)

                (singleIdx_low,) = np.where(
                    binary_choose_low > binary_fraction_low)
                (binaryIdx_low,) = np.where(
                    binary_choose_low <= binary_fraction_low)
            else:
                raise ValueError(
                    "You have supplied a fraction outside of 0-1. Please choose a fraction between 0 and 1."
                )
        else:
            raise ValueError(
                "You have not supplied a model or a fraction. Please choose either vanHaaften or a float"
            )

        # --- if using a different binary fraction for high-mass systems
        if (binfrac_model_msort is not None) and (type(binfrac_model_msort) == str):
            if binfrac_model_msort == "vanHaaften":
                binary_fraction_high = 1 / 2.0 + 1 / \
                    4.0 * np.log10(primary_mass[highmassIdx])
                binary_choose_high = np.random.uniform(
                    0, 1.0, primary_mass[highmassIdx].size)

                (singleIdx_high,) = np.where(
                    binary_fraction_high < binary_choose_high)
                (binaryIdx_high,) = np.where(
                    binary_fraction_high >= binary_choose_high)
            else:
                raise ValueError(
                    "You have supplied a non-supported binary fraction model. Please choose vanHaaften or a float"
                )
        elif (binfrac_model_msort is not None) and (type(binfrac_model_msort) == float):
            if (binfrac_model_msort <= 1.0) & (binfrac_model_msort >= 0.0):
                binary_fraction_high = binfrac_model_msort * \
                    np.ones(primary_mass[highmassIdx].size)
                binary_choose_high = np.random.uniform(
                    0, 1.0, primary_mass[highmassIdx].size)

                (singleIdx_high,) = np.where(
                    binary_choose_high > binary_fraction_high)
                (binaryIdx_high,) = np.where(
                    binary_choose_high <= binary_fraction_high)
            else:
                raise ValueError(
                    "You have supplied a fraction outside of 0-1. Please choose a fraction between 0 and 1."
                )
        elif (binfrac_model_msort is not None):
            raise ValueError(
                "You have not supplied a model or a fraction. Please choose either vanHaaften or a float"
            )

        # --- get pertinent info
        if (binfrac_model_msort is not None):
            stars_in_binary = np.append(
                primary_mass[highmassIdx][binaryIdx_high], primary_mass[lowmassIdx][binaryIdx_low])
            stars_in_single = np.append(
                primary_mass[highmassIdx][singleIdx_high], primary_mass[lowmassIdx][singleIdx_low])
            binary_fraction = np.append(
                binary_fraction_high[binaryIdx_high], binary_fraction_low[binaryIdx_low])
            binaryIdx = np.append(
                highmassIdx[binaryIdx_high], lowmassIdx[binaryIdx_low])
        else:
            stars_in_binary = primary_mass[lowmassIdx][binaryIdx_low]
            stars_in_single = primary_mass[lowmassIdx][singleIdx_low]
            binary_fraction = binary_fraction_low[binaryIdx_low]
            binaryIdx = lowmassIdx[binaryIdx_low]

        return (
            stars_in_binary,
            stars_in_single,
            binary_fraction,
            binaryIdx,
        )

    def sample_porb(self, mass1, mass2, rad1, rad2, porb_model, porb_max=None, size=None, **kwargs):
        """Sample the orbital period according to the user-specified model

        Parameters
        ----------
        mass1 : array
            primary masses
        mass2 : array
            secondary masses
        rad1 : array
            radii of the primaries.
        rad2 : array
            radii of the secondaries
        porb_model : string
            selects which model to sample orbital periods, choices include:
            log_uniform : semi-major axis flat in log space from RRLO < 0.5 up to 1e5 Rsun according to
            `Abt (1983) <http://adsabs.harvard.edu/abs/1983ARA%26A..21..343A>`_
            and consistent with Dominik+2012,2013
            and then converted to orbital period in days using Kepler III
            sana12 : power law orbital period between 0.15 < log(P/day) < 5.5 following
            `Sana+2012 <https://ui.adsabs.harvard.edu/abs/2012Sci...337..444S/abstract>_`
            renzo19 : power law orbital period for m1 > 15Msun binaries from
            `Sana+2012 <https://ui.adsabs.harvard.edu/abs/2012Sci...337..444S/abstract>_`
            following the implementation of
            `Renzo+2019 <https://ui.adsabs.harvard.edu/abs/2019A%26A...624A..66R/abstract>_`
            and flat in log otherwise
            raghavan10 : log normal orbital periods in days with mean_logP = 4.9
            and sigma_logP = 2.3 between 0 < log10(P/day) < 9  following
            `Raghavan+2010 <https://ui.adsabs.harvard.edu/abs/2010ApJS..190....1R/abstract>_`
            moe19 : log normal orbital periods in days with mean_logP = 4.9
            and sigma_logP = 2.3 between 0 < log10(P/day) < 9  following
            `Raghavan+2010 <https://ui.adsabs.harvard.edu/abs/2010ApJS..190....1R/abstract>_`
            but with different close binary fractions following
            `Moe+2019 <https://ui.adsabs.harvard.edu/abs/2019ApJ...875...61M/abstract>_`
        met : float
            metallicity of the population

        Returns
        -------
        porb : array
            orbital period with array size equalling array size
            of mass1 and mass2 in units of days
        aRL_over_a: array
            ratio of radius where RL overflow starts to the sampled seperation
            used to truncate the eccentricitiy distribution
        """

        # First we need to compute where RL overflow starts.  We truncate the lower-bound
        # of the period distribution there
        q = mass2 / mass1
        RL_fac = (0.49 * q ** (2.0 / 3.0)) / (
            0.6 * q ** (2.0 / 3.0) + np.log(1 + q ** 1.0 / 3.0)
        )

        q2 = mass1 / mass2
        RL_fac2 = (0.49 * q2 ** (2.0 / 3.0)) / (
            0.6 * q2 ** (2.0 / 3.0) + np.log(1 + q2 ** 1.0 / 3.0)
        )

        # include the factor for the eccentricity
        RL_max = 2 * rad1 / RL_fac
        (ind_switch,) = np.where(RL_max < 2 * rad2 / RL_fac2)
        if len(ind_switch) >= 1:
            RL_max[ind_switch] = 2 * rad2[ind_switch] / RL_fac2[ind_switch]

        # Can either sample the porb first and truncate the eccentricities at RL overflow
        # or sample the eccentricities first and truncate a(1-e) at RL overflow
        #
        # If we haven't sampled the eccentricities, then the minimum semi-major axis is at
        # RL overflow
        #
        # If we have, then the minimum pericenter is set to RL overflow
        a_min = RL_max

        if porb_model == "log_uniform":
            if porb_max is None:
                a_0 = np.random.uniform(np.log(a_min), np.log(1e5), size)
            else:
                # If in CMC, only sample binaries as wide as the local hard/soft boundary
                a_max = utils.a_from_p(porb_max,mass1,mass2)
                a_max[a_max < a_min] = a_min[a_max < a_min]
                a_0 = np.random.uniform(np.log(a_min), np.log(a_max), size)

            # convert out of log space
            a_0 = np.exp(a_0)
            aRL_over_a = a_min/a_0

            # convert to au
            rsun_au = 0.00465047
            a_0 = a_0 * rsun_au

            # convert to orbital period in years
            yr_day = 365.24
            porb_yr = ((a_0 ** 3.0) / (mass1 + mass2)) ** 0.5
            porb = porb_yr * yr_day
        elif porb_model == "sana12":
            # Same here: if using CMC, set the maximum porb to the smaller of either the
            # hard/soft boundary or 5.5 (from Sana paper)
            if porb_max is None:
                log10_porb_max = 5.5
            else:
                log10_porb_max = np.minimum(5.5,np.log10(porb_max))

            # Use the lower limit from the Sana12 distribution, unless this means the binaries are sampled at RL overflow. If so,
            # change the lower limit to a_min

            log10_porb_min = np.array([0.15]*len(a_min))
            RL_porb = utils.p_from_a(a_min,mass1,mass2)
            log10_RL_porb = np.log10(RL_porb)
            log10_porb_min[log10_porb_min <  log10_RL_porb] = log10_RL_porb[log10_porb_min < log10_RL_porb]

            porb = 10 ** utils.rndm(a=log10_porb_min, b=log10_porb_max, g=-0.55, size=size)
            aRL_over_a = a_min / utils.a_from_p(porb,mass1,mass2)

        elif porb_model == "renzo19":
            # Same here: if using CMC, set the maximum porb to the smaller of either the
            # hard/soft boundary or 5.5 (from Sana paper)
            if porb_max is None:
                log10_porb_max = 5.5

            else:
                log10_porb_max = np.minimum(5.5,np.log10(porb_max))

            # Use the lower limit from the Sana12 distribution, unless this means the binaries are sampled at RL overflow. If so,
            # change the lower limit to a_min
            log10_porb_min = np.array([0.15]*len(a_min))
            RL_porb = utils.p_from_a(a_min,mass1,mass2)
            log10_RL_porb = np.log10(RL_porb)
            log10_porb_min[log10_porb_min <  log10_RL_porb] = log10_RL_porb[log10_porb_min < log10_RL_porb]

            porb = 10 ** (np.random.uniform(log10_porb_min, log10_porb_max, size))
            (ind_massive,) = np.where(mass1 > 15)

            if type(log10_porb_max) != float:
                log10_porb_max = log10_porb_max[ind_massive]
                log10_porb_min = log10_porb_min[ind_massive]

            porb[ind_massive] = 10 ** utils.rndm(
                a=log10_porb_min[ind_massive], b=log10_porb_max, g=-0.55, size=len(ind_massive))
            aRL_over_a = a_min / utils.a_from_p(porb,mass1,mass2)

        elif porb_model == "raghavan10":
            import scipy
            # Same here: if using CMC, set the maximum porb to the smaller of either the
            # hard/soft boundary or 5.5 (from Sana paper)
            if porb_max is None:
                log10_porb_max = 9.0
            else:
                log10_porb_max = np.minimum(5.5, np.log10(porb_max))

            lower = 0
            upper = log10_porb_max
            mu = 4.9
            sigma = 2.3

            porb = 10 ** (scipy.stats.truncnorm.rvs(
                (lower-mu)/sigma,(upper-mu)/sigma, loc=mu, scale=sigma, size=size
                ))

            aRL_over_a = a_min / utils.a_from_p(porb,mass1,mass2)

        elif porb_model == "moe19":
            from scipy.interpolate import interp1d
            from scipy.stats import norm
            from scipy.integrate import trapz
            try:
                met = kwargs.pop('met')
            except:
                raise ValueError(
                    "You have chosen moe19 for the orbital period distribution which is a metallicity-dependent distribution. "
                    "Please specify a metallicity for the population."
                    )
            def get_logP_dist(nsamp, norm_wide, norm_close, mu=4.4, sigma=2.1):
                logP_lo_lim=0
                logP_hi_lim=9
                close_logP=4.0
                wide_logP=6.0
                neval = 500
                prob_wide = norm.pdf(np.linspace(wide_logP, logP_hi_lim, neval), loc=mu, scale=sigma)*norm_wide
                prob_close = norm.pdf(np.linspace(logP_lo_lim, close_logP, neval), loc=mu, scale=sigma)*norm_close
                slope = -(prob_close[-1] - prob_wide[0]) / (wide_logP - close_logP)
                prob_intermediate = slope * (np.linspace(close_logP, wide_logP, neval) - close_logP) + prob_close[-1]
                prob_interp_int = interp1d(np.linspace(close_logP, wide_logP, neval), prob_intermediate)

                log_p_success = []
                n_success = 0
                while n_success < nsamp:
                    logP_samp = np.random.uniform(logP_lo_lim, logP_hi_lim, nsamp*5)
                    logP_prob = np.random.uniform(0, 1, nsamp*5)

                    logP_samp_lo = logP_samp[logP_samp<close_logP]
                    logP_prob_lo = logP_prob[logP_samp<close_logP]
                    log_p_success.extend(logP_samp_lo[np.where(logP_prob_lo < norm.pdf(logP_samp_lo, loc=mu, scale=sigma)*norm_close)])

                    logP_samp_int = logP_samp[(logP_samp>=close_logP) & (logP_samp<wide_logP)]
                    logP_prob_int = logP_prob[(logP_samp>=close_logP) & (logP_samp<wide_logP)]
                    log_p_success.extend(logP_samp_int[np.where(logP_prob_int < prob_interp_int(logP_samp_int))])

                    logP_samp_hi = logP_samp[(logP_samp>=wide_logP)]
                    logP_prob_hi = logP_prob[(logP_samp>=wide_logP)]

                    log_p_success.extend(logP_samp_hi[np.where(logP_prob_hi < norm.pdf(logP_samp_hi, loc=mu, scale=sigma)*norm_wide)])

                    n_success = len(log_p_success)
                log_p_success = np.array(log_p_success)[np.random.randint(0,n_success,nsamp)]
                return log_p_success
            norm_wide, norm_close = utils.get_porb_norm(met)
            logP_dist = get_logP_dist(size, norm_wide, norm_close)
            logP_dist = logP_dist[np.random.randint(0, len(logP_dist), size)]
            porb = 10**logP_dist
            aRL_over_a = a_min / utils.a_from_p(porb,mass1,mass2)

        else:
            raise ValueError(
                "You have supplied a non-supported model; Please choose either log_flat, sana12, renzo19, raghavan10, or moe19"
            )
        return porb, aRL_over_a

    def sample_ecc(self, aRL_over_a, ecc_model="sana12", size=None):
        """Sample the eccentricity according to a user specified model

        Parameters
        ----------
        ecc_model : string
            'thermal' samples from a  thermal eccentricity distribution following
            `Heggie (1975) <http://adsabs.harvard.edu/abs/1975MNRAS.173..729H>`_
            'uniform' samples from a uniform eccentricity distribution
            'sana12' samples from the eccentricity distribution from
            `Sana+2012 <https://ui.adsabs.harvard.edu/abs/2012Sci...337..444S/abstract>_`
            'circular' assumes zero eccentricity for all systems
            DEFAULT = 'sana12'

        aRL_over_a : ratio of the minimum seperation (where RL overflow starts)
            to the sampled semi-major axis.  Use this to truncate the eccentricitiy

        size : int, optional
            number of eccentricities to sample
            this is set in cosmic-pop call as Nstep

        Returns
        -------
        ecc : array
            array of sampled eccentricities with size=size
        """

        # if we sampled the periods first, we need to truncate the eccentricities
        # to avoid RL overflow/collision at pericenter
        e_max = 1.0 - aRL_over_a

        if ecc_model == "thermal":
            a_0 = np.random.uniform(0.0, e_max**2, size)
            ecc = a_0 ** 0.5
            return ecc

        elif ecc_model == "uniform":
            ecc = np.random.uniform(0.0, e_max, size)
            return ecc

        elif ecc_model == "sana12":
            sana_max = np.array([0.9]*len(e_max))
            max_e = np.minimum(e_max, sana_max)
            ecc = utils.rndm(a=0.001, b=max_e, g=-0.45, size=size)

            return ecc

        elif ecc_model == "circular":
            ecc = np.zeros(size)
            return ecc

        else:
            raise ValueError("You have specified an unsupported model. Please choose from thermal, "
                             "uniform, sana12, or circular")

    def sample_SFH(self, SF_start=13700.0, SF_duration=0.0, met=0.02, size=None):
        """Sample an evolution time for each binary based on a user-specified
        time at the start of star formation and the duration of star formation.
        The default is a burst of star formation 13,700 Myr in the past.

        Parameters
        ----------
        SF_start : float
            Time in the past when star formation initiates in Myr
        SF_duration : float
            Duration of constant star formation beginning from SF_Start in Myr
        met : float
            metallicity of the population [Z_sun = 0.02]
            Default: 0.02
        size : int, optional
            number of evolution times to sample
            NOTE: this is set in cosmic-pop call as Nstep

        Returns
        -------
        tphys : array
            array of evolution times of size=size
        metallicity : array
            array of metallicities
        """

        if (SF_start > 0.0) & (SF_duration >= 0.0):
            tphys = np.random.uniform(SF_start - SF_duration, SF_start, size)
            metallicity = np.ones(size) * met
            return tphys, metallicity
        else:
            raise ValueError(
                'SF_start and SF_duration must be positive and SF_start must be greater than 0.0')

    def set_kstar(self, mass):
        """Initialize stellar types according to BSE classification
        kstar=1 if M>=0.7 Msun; kstar=0 if M<0.7 Msun

        Parameters
        ----------
        mass : array
            array of masses

        Returns
        -------
        kstar : array
            array of initial stellar types
        """

        kstar = np.zeros(mass.size)
        low_cutoff = 0.7
        lowIdx = np.where(mass < low_cutoff)[0]
        hiIdx = np.where(mass >= low_cutoff)[0]

        kstar[lowIdx] = 0
        kstar[hiIdx] = 1

        return kstar

    def set_reff(self, mass, metallicity, zsun=0.02):
        """
        Better way to set the radii from BSE, by calling it directly

        takes masses and metallicities, and returns the radii

        Note that the BSE function is hard-coded to go through arrays
        of length 10^5.  If your masses are more than that, you'll
        need to divide it into chunks
        """

        max_array_size = 100000
        total_length = len(mass)
        radii = np.zeros(total_length)

        _evolvebin.metvars.zsun = zsun

        idx = 0
        while total_length > max_array_size:
            ## cycle through the masses max_array_size number at a time
            temp_mass = mass[idx*max_array_size:(idx+1)*max_array_size]

            temp_radii = _evolvebin.compute_r(temp_mass,metallicity,max_array_size)

            ## put these in the radii array
            radii[idx*max_array_size:(idx+1)*max_array_size] = temp_radii

            total_length -= max_array_size
            idx += 1

        length_remaining = total_length

        ## if smaller than 10^5, need to pad out the array
        temp_mass = np.zeros(max_array_size)
        temp_mass[:length_remaining] = mass[-length_remaining:]

        temp_radii = _evolvebin.compute_r(temp_mass,metallicity,length_remaining)

        #finish up the array
        radii[-length_remaining:] = temp_radii[:length_remaining]

        return radii


import cv2 as cv2
from cv2 import cv2
import numpy as np
import math as m

frame_height = 1000
frame_width = 1000

number_of_points = 25
center = (frame_width/2, frame_height/2)
radius = 400

font = cv2.FONT_HERSHEY_SIMPLEX
green = (255, 0, 0)
blue = (0, 255, 0)
red = (0, 0, 255)
white = (255, 255, 255)
black = (0, 0, 0)

iterations = 0

number_of_rows = number_of_points - 1
path = r'C:\Users\spenc\Downloads\Anchor.png'   #image path
img = cv2.imread(path)
gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

dictionary_of_pixels = {}
dictionary_of_strings = {}

def resizeImage(img_to_resize):
    img_height = len(img_to_resize)
    img_width = len(img_to_resize[0])
    resized_image = img_to_resize

    if img_height < frame_height:
        height_buffer = np.full(    ( int((frame_height - img_height)/2) , img_width), 255, dtype=np.uint8)
        resized_image = np.concatenate((height_buffer, img_to_resize, height_buffer), axis=0)
        img_height = len(resized_image)
        #print("image height is " + str(img_height))
    img_height = len(resized_image)
    if img_width < frame_width:
        #print("this new image height is: " + str(img_height))
        width_buffer = np.full(     (img_height,  int((frame_width - img_width)/2) ), 255, dtype=np.uint8)
        #print(str(len(width_buffer)) + " is the height of buffer")
        #print(str(len(width_buffer[0])) + " is the length of buffer")

        resized_image = np.concatenate((width_buffer, resized_image, width_buffer), axis=1)
    return(resized_image)

def cutCircle(img_to_cut):    #crops image into circle, removes all points besides ones inside circle
    global radius           #radius of circle
    global frame_height, frame_width    #height and width of string img
    pixels_to_keep = []
    pixels_to_remove = []
    img_height = len(img_to_cut)  #width and height of image in pixels
    img_width = len(img_to_cut[0])
    this_center_x = img_width / 2
    this_center_y = img_height / 2
    for y in range(len(img_to_cut)):
        for x in range(len(img_to_cut[y])):
            if distance((x,y), (this_center_x, this_center_y)) < radius:
                pixels_to_keep.append((x,y))
            else:
                pixels_to_remove.append((x,y))
    return(pixels_to_keep, pixels_to_remove)

def blackOutPixels(these_pixels_to_remove, this_img):
    copy_of_image = this_img
    for p in these_pixels_to_remove:
        x, y = p
        if y < len(copy_of_image) and x < len(copy_of_image[0]):
            copy_of_image[y, x] = 0
    return copy_of_image

def invertThisImage(img_to_invert):
    t_inverted_img = (255 - img_to_invert)
    return t_inverted_img

def distance(p1, p2): #absolute distance between points p1 and p2
    x1, y1 = p1
    x2, y2 = p2
    distance = m.sqrt(  (x1-x2)**2  +   (y1-y2)**2  )
    return(distance)

def nextColor():
    global iterations
    bgr = hsv2bgr((100 / number_of_points * iterations, 180, 180))
    iterations = iterations + 1
    return bgr

def cvWindows():
    global frame, final_strings
    frame = np.zeros(shape=(frame_height, frame_width, 3), dtype=np.uint8)
    final_strings = frame
def connectPoints(number_of_points):

    if number_of_points > 1:
        for i in range(number_of_points - 1):
           list_of_strings.append(makeString(number_of_points, i + 1, hsv2bgr((20 * i, 180, 180)), frame))
        number_of_points = number_of_points - 1
        connectPoints(number_of_points)

def pixelsInaLine(this_string):

    l_m, l_b = stringToLine(this_string)
    #print("m and b are: " + str(l_m) +" "+ str(l_b))
    pixels_in_line = []
    if l_m is "vert":
        x = l_b
        for y in range(frame_height):
            if distance((x,y), center) <= radius:
                pixel = str((x,y))
                if pixel in dictionary_of_pixels.keys():
                    list_of_lines = dictionary_of_pixels[pixel]
                    list_of_lines.append(this_string)
                    dictionary_of_pixels[pixel] = list_of_lines
                else:
                    list_of_lines = [this_string]
                    dictionary_of_pixels[pixel] = list_of_lines
    else:
        for x in range(frame_width):
            y = int(l_m * x + l_b)
            if distance((x,y), center) <= radius:
                pixel = str((x, y))
                if pixel in dictionary_of_pixels.keys():
                    list_of_lines = dictionary_of_pixels[pixel]
                    list_of_lines.append(this_string)
                    dictionary_of_pixels[pixel] = list_of_lines
                    #print("adding line..." + pixel + " " + str(dictionary_of_pixels[pixel]))
                else:
                    list_of_lines = [this_string]
                    dictionary_of_pixels[pixel] = list_of_lines

def createIntersections(list_of_strings):
    #list_of_intersections = []
    pass

def makeString(id1, id2, color, image_to_draw):
    p1 = pointToCoord(id1)
    p2 = pointToCoord(id2)
    cv2.line(image_to_draw, p1, p2, color, 1)
    if id1 < id2:
        t_string = (id1, id2)
    else:
        t_string = (id2, id1)
    return t_string

def intersection(string1, string2):
    m1, b1= stringToLine(string1)
    m2, b2= stringToLine(string2)
    if m1 != m2:
        if m1 == "vert":
            intersection_x = b1
            intersection_y = intersection_x * m2 + b2
        elif m2 == "vert":
            intersection_x = b2
            intersection_y = intersection_x * m1 + b1
        elif ((m1 - m2) != 0):
            intersection_x = (b2 - b1) / (m1 - m2)
            intersection_y = intersection_x * m1 + b1
    else:
        intersection_x = 0
        intersection_y = 0
    return (int(intersection_x), int(intersection_y))

def stringToLine(string):
    x1, y1 = pointToCoord(string[0])
    x2, y2 = pointToCoord(string[1])
    if x1 != x2:
        m = (y1 - y2) / (x1 - x2)
        b = y1 - m * x1
    else:
        m = "vert"
        b = x1
    return(m, b)

def pointToCoord(pointID):
    global center
    center_x, center_y = center
    angle = pointID * 2 * m.pi / number_of_points
    point_x = int(m.cos(angle) * radius + center_x)
    point_y = int(m.sin(angle) * radius + center_y)
    cv2.circle(frame, (point_x, point_y), 1, white, -1)
    cv2.putText(frame, str(pointID), (point_x, point_y), font, 0.5, white)

    return(point_x, point_y)

def hsv2bgr(h_s_v):

    #bgr_red = np.uint8([[[0, 0, 255]]])
    #hsv_red = cv2.cvtColor(bgr_red, cv2.COLOR_BGR2HSV)

    hsv_space = np.uint8([[h_s_v]])
    bgr_space = cv2.cvtColor(hsv_space, cv2.COLOR_HSV2BGR)
    b = int(bgr_space[0, 0, 0])
    g = int(bgr_space[0, 0, 1])
    r = int(bgr_space[0, 0, 2])
    return (b, g, r)

def drawPixels(dictionary_to_draw, image):
    for pixel in dictionary_to_draw:
        #print(pixel)
        x, y = pixel.split((", "))
        _, x = x.split("(")
        y, _ = y.split(")")
        x = int(x)
        y = int(y)
        #print(x, y)
        total_value_for_this_pixel = 0
        for t_string in dictionary_to_draw[pixel]:
            value_for_this_string = dictionary_of_strings[str(t_string)]
            total_value_for_this_pixel += value_for_this_string

        image[y, x] = total_value_for_this_pixel
        '''
        strings_for_pixel = dictionary_of_pixels[pixel]
        for t_string in strings_for_pixel:
            if t_string in dictionary_of_strings.keys():
                number_of_strings = dictionary_of_strings[t_string]
                number_of_strings += 1
                dictionary_of_strings[t_string] = number_of_strings
            else:
                number_of_strings = 1
                dictionary_of_strings[t_string] = number_of_strings
        '''

def assignValuesToStrings(inverted_image):
    for y in range(len(inverted_image)):
        for x in range(len(inverted_image[y])):
            pixel = str((x, y))
            if pixel in dictionary_of_pixels.keys():
                strings_for_pixel = dictionary_of_pixels[pixel]
                number_of_strings = len(strings_for_pixel)
                for individual_string in strings_for_pixel:
                    t_string = str(individual_string)
                    if t_string in dictionary_of_strings.keys():
                        number_of_strings = dictionary_of_strings[t_string]
                        number_of_strings += 1 / ((inverted_image[y, x] + 1))
                        dictionary_of_strings[t_string] = number_of_strings
                    else:
                        number_of_strings = 1
                        dictionary_of_strings[t_string] = number_of_strings
            else:
                #print("false, because " + pixel)
                pass

cvWindows()
list_of_strings = []
coords = []
cv2.rectangle(frame, (0, 0), (frame_width, frame_height), (20, 20, 20), -1)

connectPoints(number_of_points)
resized_image = resizeImage(gray)

pixels_to_keep, pixels_to_remove = cutCircle(resized_image)
cut_out_image = blackOutPixels(pixels_to_remove, resized_image)

#print(list_of_strings)

for string in list_of_strings:
    pixelsInaLine(string)
#print(dictionary_of_pixels)
inverted_image = invertThisImage(cut_out_image)

assignValuesToStrings(inverted_image)

drawPixels(dictionary_of_pixels, cut_out_image)
#print("this stuff right here is: " + str(len(dictionary_of_pixels["(500, 500)"])))
for t_string in dictionary_of_strings:
    value_for_this_string = dictionary_of_strings[t_string]
    p1, p2 = string
    makeString(p1, p2, (value_for_this_string, value_for_this_string, value_for_this_string), final_strings)

print("dictionary of strings:")
print(dictionary_of_strings)

while(1):
    #angle_x = cv2.getTrackbarPos("rotation about x", "frame")
    cv2.imshow("resized", resized_image)
    cv2.imshow("cut out", cut_out_image)
    cv2.imshow('strings', frame)
    cv2.imshow("final_strings", final_strings)

    key_value = cv2.waitKey(30)
    if key_value == ord('q'):
        break


from django.conf.urls import patterns, include, url
from rest_framework import routers
from lettercounter.counters import views

import settings

# Uncomment the next two lines to enable the admin:
# from django.contrib import admin
# admin.autodiscover()

url_patterns = [
    '',
    url(r'^api/counts/$', views.CountList.as_view()),
    url(r'^api/counts/(?P<pk>[A-Za-z]+)/$', views.CountDetail.as_view()),
    url(r'^api/sentences/$', views.SentenceDetail.as_view()),
    url(r'^api/stats/$', views.StatsDetail.as_view()),
]

if settings.DEBUG:
    url_patterns.append(
        url(
            r'^$', 'django.contrib.staticfiles.views.serve', kwargs={
                'path': 'index.html', 'document_root': settings.STATIC_ROOT}),
        )

url_patterns.append(
    url(r'^(?P<path>.*)$', 'django.contrib.staticfiles.views.serve', {
        'document_root': settings.STATIC_ROOT, 'show_indexes': False
    })
)

urlpatterns = patterns(*url_patterns)


# Itertools
import itertools

# Infinite Counting
for x in itertools.count(50, 5): # count as 50, 55, 60... infinitely
    print(x)
    if x == 80: # stopper condition
        break

# Infinite Cycling
counter = 0
for c in itertools.cycle("Apple Is so wonderful!"):
    print(c)

    counter += 1
    if counter > 50:
        break

# Infinite Cycling for number
counter = 0
for c in itertools.cycle([1,2,3,4,5]):
    print(c)

    counter += 1
    if counter > 50:
        break

# Infinite Repeating
counter = 0
for r in itertools.repeat("moon!"):
    print(r)

    counter += 1
    if counter > 50:
        break

13.Learning Python Libraries/01.Learning Python Standard Library - KH/02.Python Math/05_itertools_processes.py


"""
General functions for managing and creating virtual environments.
"""
import os
import warnings

from .renv import create_renv_env
from .conda import get_conda_activate_command, create_conda_env
from ..utils.utils import build_model_identifier

def get_model_env_types(path):
    """Determine whether a model specifies a conda or renv environment. The presence of
    a conda environment is defined by the existence of the file path/environment.yml,
    or the file path/renv.lock for a renv environment.

    Parameters
    ----------
    path : str
        Path to model directory

    Returns
    -------
    dict
        {env_type: env_found} dictionary, where env_type is the type of environment and
        env_found is True if an environment of that type was discovered and False
        otherwise.
    """
    if os.path.exists(f"{path}/environment.yml"):
        conda = True
    else:
        conda = False

    if os.path.exists(f"{path}/renv.lock"):
        renv = True
    else:
        renv = False

    return {"conda": conda, "renv": renv}

def create_env(model_path, model_id, model_version, capture_output=False):
    """Create all the environments defined in a model directory.

    Parameters
    ----------
    model_path : str
        Path to model directory
    model_id : int
        ID for the model
    model_version : str
        Version of the model
    capture_output : bool, optional
        Passed to subprocess.run, by default False

    Returns
    -------
    str
       Command needed to activate the model environment. If both conda and renv
       environments specified this will be the command to activate the conda environment
       only.
    """
    env_types = get_model_env_types(model_path)
    env_cmd = None

    if env_types["renv"] and env_types["conda"]:
        warnings.warn(
            "Both conda and renv environment detected - "
            "conda environment will be given priority."
        )

    if env_types["renv"]:
        conda_env = create_renv_env(model_path, capture_output=capture_output)
        # create_renv_env returns conda environment with appropriate R version
        #  (as well as setting up renv), if set in config
        if conda_env:
            env_cmd = get_conda_activate_command(conda_env)

    if env_types["conda"]:
        env_name = build_model_identifier(model_id, model_version)
        create_conda_env(
            env_name,
            env_file=f"{model_path}/environment.yml",
            capture_output=capture_output,
        )
        env_cmd = get_conda_activate_command(env_name)

    return env_cmd


from flask import Flask, render_template, request, redirect, flash
# from controllers import users
from controllers import bases
app = Flask(__name__)
app.secret_key = b'_5#y2L"F4Q8z\n\xec]/'
@app.get('/')
def home():
    return render_template('index.html')

@app.post('/guardar')
def guardarUsuario():
    print('validar')
    data = request.form

    print(data['contra'])

    if( data.getlist('gridCheck')):

        datos = bases.listaBases(data)
        print(datos)
        return render_template('index.html',datosDB = datos)
    else:

        bases.crear(data)
        return redirect('/')
app.run(debug=True)


import sys
n, r, c = map(int, sys.stdin.readline().split())
target = 0

while True:
    h = 2**(2*(n-1))
    if r >=2**(n-1):
        r -= 2**(n-1)
        target += 2*h
    if c >=2**(n-1):
        c -= 2**(n-1)
        target += h
    n -= 1
    if r == 0 and c == 0:
        break
print(target)

import sys
n = int(sys.stdin.readline())
rem_cnt = sys.stdin.readline()
rem = list(map(int, sys.stdin.readline().split()))
arr = [x for x in range(10) if x not in rem]
a = b = 0
n2 = n - 1
if int(rem_cnt) == 0:
    pass
while True:
    flag = 1
    n2 += 1
    for i in str(n2):
        if int(i) not in arr:
            flag = 0
            break
    if flag == 1:
        a = len(str(n2)) + (n2-n)
        break
n2 = n
while True:
    flag = 1
    n2 -= 1
    for i in str(n2):
        if int(i) not in arr:
            flag = 0
            break
    if flag == 1 or n2 < 0:
        b = len(str(n2)) + (n-n2)
        break
print(min(a, b))

import sys
from collections import deque
def bfs(graph, root):
    visited = []
    queue = deque([root])
    cnt = 0
    while queue:
        n = queue.popleft()
        if n not in visited:
            visited.append(n)
            if n in graph:
                tmp = list(set(graph[n])-set(visited))
                tmp.sort()
                queue += tmp
        cnt += 1
    return cnt

graph = {}
n, m = map(int, sys.stdin.readline().split())
for i in range(m):
    n1, n2 = map(int, sys.stdin.readline().split())
    if n1 not in graph:
        graph[n1] = [n2]
    elif n2 not in graph:
        graph[n2] = [n1]

    if n2 not in graph:
        graph[n2] = [n1]
    elif n1 not in graph:
        graph[n1] = [n2]
ans = []
for i in graph:
    ans.append(bfs(graph, i))
print(min(ans))

import sys
from math import inf
n = int(sys.stdin.readline())
m = int(sys.stdin.readline())

graph = [[inf]*n for x in range(n)]
for i in range(m):
    a, b, c = map(int, sys.stdin.readline().split())
    graph[a-1][b-1] = min(graph[a-1][b-1], c)

for i in range(n):
    graph[i][i] = 0
    for j in range(n):
        for k in range(n):
            graph[j][k] = min(graph[j][k], graph[j][i]+graph[i][k])

for i in graph:
    for j in i:
        if j == inf:
            j = 0
        print(j, end=" ")
    print()

import sys

a, b = map(int, sys.stdin.readline().split())
graph = [[a]*a for x in range(a)]
for i in range(b):
    n1, n2 = map(int, sys.stdin.readline().split())
    graph[n1-1][n2-1] = min(graph[n1-1][n2-1], 1)
    graph[n2-1][n1-1] = min(graph[n2-1][n1-1], 1)
for k in range(a):
    graph[k][k] = 0
for k in range(a):
    for i in range(a):
        for j in range(a):
            graph[i][j] = min(graph[i][j], graph[i][k]+graph[k][j])
minarr = []
for i in graph:
    if sum(i) not in graph:
        minarr.append(sum(i))
    else:
        minarr.append(sum(i)+1)
print(minarr.index(min(minarr))+1)


import pandas as pd
import streamlit as st
import numpy as np
import time

data = pd.read_csv('medic.csv')

a = [1,2,3,4,5,6,7,8]
n = np.array(a)
nd = np.array([[1,4,6], [5,7,8]])

dic = {
    "name" : ['aadm', 'bola', 'jide'],
    "age": [12,15,89]
}

st.dataframe(data)
# st.table(data)

@st.cache
def run_time(delay=8):
    time.sleep(delay)
    return time.time()

if st.checkbox('1'):
    st.write(run_time())

if st.checkbox('2'):
    st.write(run_time())


import matplotlib.pyplot as plt

x = ['A', 'B', 'C', 'D', 'E']
y = [3, 7, 2, 5, 9]

plt.bar(x, y)
plt.title('Bar Plot')
plt.xlabel('Categories')
plt.ylabel('Values')
plt.show()


import argparse
import os
import random
import array
import multiprocessing

import yaml
import operator
import numpy as np
import pandas as pd
from scipy.optimize import minimize as neldermead

from reinforcement_learning.config import PROJECT_PATH
from reinforcement_learning.pump_gym.envs.water_gym import WaterGym

parser = argparse.ArgumentParser()
parser.add_argument('--params', default='anytownMaster',
                    type=str, help="Name of the YAML file.")
parser.add_argument('--nscenes', default=100, type=int,
                    help="Number of the scenes to generate.")
parser.add_argument('--seed', default=None, type=int,
                    help="Random seed for the optimization methods.")
parser.add_argument('--scenes', default=None, type=str,
                    help="Name of the generated scenes file.")
parser.add_argument('--result', default=None, type=str,
                    help="Name of the generated result file.")
parser.add_argument('--nproc', default=1, type=int,
                    help="Number of processes to raise.")
args = parser.parse_args()

path_to_params = os.path.join(
    PROJECT_PATH, 'pump_gym', 'experiments', 'hyperparameters', args.params+'.yaml')
with open(path_to_params, 'r') as fin:
    hparams = yaml.load(fin, Loader=yaml.Loader)

reset_orig_demands = hparams['env']['reset_original_demands']
wn_name = hparams['env']['water_network']+'_master'

if args.scenes:
    scene_file_name = args.scenes + '.csv'
else:
    scene_file_name = hparams['evaluation']['scenes'] + '.csv'

if args.result:
    result_file_name = args.result + '.csv'
else:
    result_file_name = hparams['evaluation']['result'] + '.csv'

n_scenes = args.nscenes
seed = args.seed
n_proc = args.nproc

if seed:
    random.seed(args.seed)
    np.random.seed(args.seed)
else:
    random.seed()
    np.random.seed()

if n_scenes > 10:
    verbosity = n_scenes // 10
else:
    verbosity = 1

env = WaterGym(
    wn_name=hparams['env']['water_network']+'_master',
    speed_increment=hparams['env']['speed_increment'],
    episode_len=hparams['env']['episode_length'],
    pump_groups=hparams['env']['pump_groups'],
    total_demand_low=hparams['env']['total_demand_low'],
    total_demand_high=hparams['env']['total_demand_high'],
    reset_original_demands=hparams['env']['reset_original_demands'],
    reset_original_pump_speeds=hparams['env']['reset_original_pump_speeds'],
    seed=args.seed
)

def reward_to_scipy(pump_speeds):
    # print('optimizing speed')
    # print(pump_speeds)
    """Only minimization allowed."""
    return -env.get_state_value_to_opti(pump_speeds)

def generate_scenes(reset_orig_demands, n_scenes):
    junction_ids = list(env.wn.junction_name_list)
    demand_df = pd.DataFrame(
        np.empty(shape=(n_scenes, len(junction_ids))),
        columns=junction_ids)
    if reset_orig_demands:
        for i in range(n_scenes):
            demand_df.loc[i] = env.build_demand_dict()
    else:
        for i in range(n_scenes):
            env.randomize_demands()
            demand_df.loc[i] = env.build_demand_dict()

    return demand_df

class Nelder_mead_method():
    def __init__(self):
        self.options = {'maxfev': 1000, 'xatol': .005, 'fatol': .01}

    def maximize(self, scene_id):
        if seed:
            random.seed(args.seed)
            init_guess = []
            for i in range(env.dimensions):
                init_guess.append(random.uniform(
                    env.speed_limit_low, env.speed_limit_high))
        else:
            random.seed()
            init_guess = []
            for i in range(env.dimensions):
                init_guess.append(random.uniform(
                    env.speed_limit_low, env.speed_limit_high))

        env.set_demands(scene_df.loc[scene_id])
        # init_guess = [1.2, 0.83369588, 0.87215528, 1.04711067, 1.07251507]
        options = {'maxfev': 1000, 'xatol': .005, 'fatol': .01}
        result = neldermead(
            reward_to_scipy,
            init_guess,
            method='Nelder-Mead',
            options=options)

        result_df = pd.DataFrame(
            np.empty(shape=(1, len(df_header))),
            columns=df_header)
        result_df['index'] = scene_id
        result_df['reward'] = -result.fun
        result_df['evals'] = result.nit

        for i in range(env.dimensions):
            result_df['speed_of_group'+str(i)] = result.x[i]
        return result_df

def optimize_scenes(scene_df, method=None):
    # pool        = multiprocessing.Pool(n_proc)
    # result_df   = pool.map(method, range(len(scene_df)))
    result_df = map(method, range(len(scene_df)))
    result_df = pd.concat(result_df)
    result_df.set_index('index', inplace=True)
    result_df.rename_axis(None, inplace=True)
    return result_df

scene_df = generate_scenes(reset_orig_demands, n_scenes)
scene_df.to_csv(os.path.join(
    PROJECT_PATH, 'pump_gym', 'data', scene_file_name))

df_header = ['index', 'reward', 'evals']
for i in range(env.dimensions):
    df_header.append('speed_of_group'+str(i))

nm = Nelder_mead_method()
subdf_nm = optimize_scenes(scene_df, nm.maximize)

subdf_nm.to_csv(os.path.join(
    PROJECT_PATH, 'pump_gym', 'data', result_file_name))


from pathlib import Path
import typing as T

from . import logging
from . import shell
from .toolz import *

log = logging.new_log(__name__)

@curry
@ensure_paths
def grep_file(regexes: T.Sequence[Regex], path: Path, *,
              getoutput=None, echo: bool=True,
              **grep_options):

    log.info(
        f'Performing grep action for path: {path}'
    )
    getoutput = getoutput or shell.getoutput(
        echo=echo,
    )

    options_string = shell.options_string(grep_options)

    log.debug(f'  grep options: {options_string}')

    regex_string = pipe(
        regexes,
        map(lambda r: f"-e '{r}'"),
        ' '.join,
    )
    log.debug(f'  regexes: {regex_string}')

    command = f'''
    grep {regex_string} {options_string} {path}
    '''.strip()

    log.info(f'  command: {command}')

    return getoutput(command)

egrep_file = grep_file(extended_regexp = True)
pgrep_file = grep_file(perl_regexp = True)

@curry
@ensure_paths
def recursive_grep(path_filter_f: T.Callable[[Path], bool], root: Path):
    pass


# coding: UTF-8
import time
import torch
import numpy as np
from train_eval import train, init_network
from importlib import import_module
import argparse
from utils import build_dataset, build_iterator, get_time_dif
import pandas as pd

parser = argparse.ArgumentParser(description='Chinese Text Classification')
parser.add_argument('--model', type=str, required=True, help='choose a model: Bert, ERNIE')
args = parser.parse_args()

if __name__ == '__main__':
    dataset = 'THUCNews'  # 数据集
    # 需要训练什么模型就自己给model_name命名，名字采用models包下的名字
    model_name = args.model  # bert
    x = import_module('models.' + model_name)
    config = x.Config(dataset)
    # np.random.seed(1)
    # torch.manual_seed(1)
    # torch.cuda.manual_seed_all(1)
    # torch.backends.cudnn.deterministic = True  # 保证每次结果一样

    i = 0
    start_time = time.time()
    print("Loading data...")

    # 读取数据
    data_train = pd.read_csv('THUCNews/data/data_train.csv', engine="python", encoding="utf_8_sig")
    # 划分边界
    boundary1 = int(data_train.shape[0] * i / 5.0)
    boundary2 = int(data_train.shape[0] * (i+1) / 5.0)
    # 取中间的那一份作为验证集
    dev_set = data_train[boundary1:boundary2]
    # 剩下的数据合并为训练集
    train_set = pd.merge(data_train[:boundary1], data_train[boundary2:], how="outer")
    # 将数据保存到模型定义好的路径里面去
    train_set.to_csv('THUCNews/data/train.csv', index=False, header=False)
    dev_set.to_csv('THUCNews/data/dev.csv', index=False, header=False)
    # train_data, dev_data, test_data = build_dataset(config)
    train_data, dev_data = build_dataset(config)
    train_iter = build_iterator(train_data, config)
    dev_iter = build_iterator(dev_data, config)
    # test_iter = build_iterator(test_data, config)
    time_dif = get_time_dif(start_time)
    print("Time usage:", time_dif)

    # train
    model = x.Model(config).to(config.device)
    # train(config, model, train_iter, dev_iter, test_iter)
    train(config, model, train_iter, dev_iter, save_path="THUCNews/saved_dict/"+ model_name + str(i) + '.ckpt')


import argparse
import cv2
import numpy as np
import torch
from torch.autograd import Function
from torchvision import models

class FeatureExtractor():
    """ Class for extracting activations and
    registering gradients from targetted intermediate layers """

    def __init__(self, model, target_layers):
        self.model = model
        self.target_layers = target_layers
        self.gradients = []

    def save_gradient(self, grad):
        self.gradients.append(grad)

    def __call__(self, x):
        outputs = []
        self.gradients = []
        for name, module in self.model._modules.items():
            x = module(x)
            if name in self.target_layers:
                x.register_hook(self.save_gradient)
                outputs += [x]
        return outputs, x

class ModelOutputs():
    """ Class for making a forward pass, and getting:
    1. The network output.
    2. Activations from intermeddiate targetted layers.
    3. Gradients from intermeddiate targetted layers. """

    def __init__(self, model, target_layers):
        self.model = model
        self.feature_extractor = FeatureExtractor(self.model.features, target_layers)

    def get_gradients(self):
        return self.feature_extractor.gradients

    def __call__(self, x):
        target_activations, output = self.feature_extractor(x)
        output = output.view(output.size(0), -1)
        output = self.model.classifier(output)
        return target_activations, output

def preprocess_image(img):
    means = [0.485, 0.456, 0.406]
    stds = [0.229, 0.224, 0.225]

    preprocessed_img = img.copy()[:, :, ::-1]
    for i in range(3):
        preprocessed_img[:, :, i] = preprocessed_img[:, :, i] - means[i]
        preprocessed_img[:, :, i] = preprocessed_img[:, :, i] / stds[i]
    preprocessed_img = \
        np.ascontiguousarray(np.transpose(preprocessed_img, (2, 0, 1)))
    preprocessed_img = torch.from_numpy(preprocessed_img)
    preprocessed_img.unsqueeze_(0)
    input = preprocessed_img.requires_grad_(True)
    return input

def show_cam_on_image(img, mask, name):
    heatmap = cv2.applyColorMap(np.uint8(255*mask), cv2.COLORMAP_JET)
    heatmap = np.float32(heatmap) / 255
    cam = heatmap + np.float32(img)
    cam = cam / np.max(cam)
    cv2.imwrite(name + ".jpg", np.uint8(255 * cam))

class GradCam:
    def __init__(self, model, target_layer_names, use_cuda):
        self.model = model
        self.model.eval()
        self.cuda = use_cuda
        if self.cuda:
            self.model = model.cuda()

        self.extractor = ModelOutputs(self.model, target_layer_names)

    def forward(self, input):
        return self.model(input)

    def __call__(self, input, index=-1):
        if self.cuda:
            features, output = self.extractor(input.cuda())
        else:
            features, output = self.extractor(input)

        if index == -1:
            index = np.argmax(output.cpu().data.numpy())

        synset = [l.strip() for l in open('./synset.txt').readlines()]
        print("class of interest: ", synset[index])

        one_hot = np.zeros((1, output.size()[-1]), dtype=np.float32)
        one_hot[0][index] = 1
        one_hot = torch.from_numpy(one_hot).requires_grad_(True)
        if self.cuda:
            one_hot = torch.sum(one_hot.cuda() * output)
        else:
            one_hot = torch.sum(one_hot * output)

        self.model.features.zero_grad()
        self.model.classifier.zero_grad()
        one_hot.backward(retain_graph=True)

        grads_val = self.extractor.get_gradients()[-1].cpu().data.numpy()

        target = features[-1]
        target = target.cpu().data.numpy()[0, :]

        # XGrad_CAM
        X_weights = np.sum(grads_val[0, :] * target, axis=(1, 2))
        X_weights = X_weights / (np.sum(target, axis=(1, 2)) + 1e-6)
        # Grad_CAM
        weights = np.mean(grads_val, axis=(2, 3))[0, :]

        X_cam = np.zeros(target.shape[1:], dtype=np.float32)
        cam = np.zeros(target.shape[1:], dtype=np.float32)

        for i, w in enumerate(weights):
            cam += w * target[i, :, :]
            X_cam += X_weights[i] * target[i, :, :]

        cam = np.maximum(cam, 0)
        cam = cv2.resize(cam, (224, 224))
        cam = cam - np.min(cam)
        cam = cam / np.max(cam)
        X_cam = np.maximum(X_cam, 0)
        X_cam = cv2.resize(X_cam, (224, 224))
        X_cam = X_cam - np.min(X_cam)
        X_cam = X_cam / np.max(X_cam)
        return cam, X_cam

def get_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('--use-cuda', action='store_true', default=False,
                        help='Use NVIDIA GPU acceleration')
    parser.add_argument('--image-path', type=str, default='./examples/both.png',
                        help='Input image path')
    parser.add_argument('--target-index', type=int, default= -1,
                        help='class of interest')
    args = parser.parse_args()
    args.use_cuda = args.use_cuda and torch.cuda.is_available()
    if args.use_cuda:
        print("Using GPU for acceleration")
    else:
        print("Using CPU for computation")

    return args

if __name__ == '__main__':
    """ python XGrad-cam.py --image-path <path_to_image> --use-cuda"""

    args = get_args()

    grad_cam = GradCam(model=models.vgg16(pretrained=True), \
                       target_layer_names=["30"], use_cuda=args.use_cuda)

    img = cv2.imread(args.image_path, 1)
    img = np.float32(cv2.resize(img, (224, 224))) / 255
    input = preprocess_image(img)

    # If -1, returns the map for the highest scoring category.
    # Otherwise, targets the requested index.
    target_index = args.target_index
    [cam, X_cam] = grad_cam(input, target_index)

    show_cam_on_image(img, cam, 'cam')
    show_cam_on_image(img, X_cam, 'X_cam')


import lxml
import lxml.etree
from collada.util import bytes, BytesIO

from pkg_resources import resource_string

# get a copy of the XML schema
# resource_string returns bytes so decode into string
COLLADA_SCHEMA_1_4_1 = str(resource_string(
    'collada',
    'resources/schema-1.4.1.xml').decode('utf-8'))

XML_XSD = str(resource_string(
    'collada',
    'resources/xsd.xml').decode('utf-8'))

class ColladaResolver(lxml.etree.Resolver):
    """COLLADA XML Resolver. If a known URL referenced
    from the COLLADA spec is resolved, a cached local
    copy is returned instead of initiating a network
    request"""

    def resolve(self, url, id, context):
        """Currently Resolves:
         * http://www.w3.org/2001/03/xml.xsd
        """
        if url == 'http://www.w3.org/2001/03/xml.xsd':
            return self.resolve_string(XML_XSD, context)
        else:
            return None

class ColladaValidator(object):
    """Validates a collada lxml document"""

    def __init__(self):
        """Initializes the validator"""
        self.COLLADA_SCHEMA_1_4_1_DOC = None
        self._COLLADA_SCHEMA_1_4_1_INSTANCE = None

    def _getColladaSchemaInstance(self):
        if self._COLLADA_SCHEMA_1_4_1_INSTANCE is None:
            self._parser = lxml.etree.XMLParser()
            self._parser.resolvers.add(ColladaResolver())
            self.COLLADA_SCHEMA_1_4_1_DOC = lxml.etree.parse(
                BytesIO(bytes(COLLADA_SCHEMA_1_4_1, encoding='utf-8')),
                self._parser)
            self._COLLADA_SCHEMA_1_4_1_INSTANCE = lxml.etree.XMLSchema(
                self.COLLADA_SCHEMA_1_4_1_DOC)
        return self._COLLADA_SCHEMA_1_4_1_INSTANCE

    COLLADA_SCHEMA_1_4_1_INSTANCE = property(_getColladaSchemaInstance)
    """An instance of lxml.XMLSchema that can be used to validate"""

    def validate(self, *args, **kwargs):
        """A wrapper for lxml.XMLSchema.validate"""
        return self.COLLADA_SCHEMA_1_4_1_INSTANCE.validate(*args, **kwargs)


'''
This file produces a heatmap based on data which is indexed by zipcode.
'''

#Source used to fix bug with matplotlib:
#https://stackoverflow.com/questions/27147300/matplotlib-tcl-asyncdelete-async-handler-deleted-by-the-wrong-thread

import pandas as pd
import geopandas as gpd
import matplotlib
matplotlib.use('Agg')
from matplotlib import pyplot as plt

def plot_heatmap(data, one):
    '''
    Creates heatmap for the data requested by the user

    Inputs:
        data (pandas Series): Contains data for the quantifier requested
            by the user, indexed by zip code
        one (bool): Saves as the heatmap for the first variable when
            one is True, and for the second variable when one is False

    Saves the heatmap as a .png
    '''
    shapefile = './res/Neighborhoods/geo_export_a131bb4a-42f9-4af9-bd33-aedec365203f.shp'
    chicago = gpd.read_file(shapefile)
    new_chicago = chicago.set_index('zip', drop=False)

    merged = pd.merge(left = new_chicago, right = data, on = new_chicago.index, how = 'inner')

    data_name = data.name
    fig, ax = plt.subplots(1, 1)
    merged.plot(column=data_name, ax=ax, legend=True, cmap='OrRd')
    # Control which heatmap is made
    if one:
        plt.savefig('static/var1.png')
    else:
        plt.savefig('static/var2.png')
    plt.close(fig)


#Import required objects from the SQLAlchemy library
from sqlalchemy import create_engine, MetaData, Table, Column, Integer, String, Float, ForeignKey

#This method sets up the database
def db_setup():
  #Creates an engine (interface) to the given database
  db = create_engine("sqlite:///database.sqlite3")

  #Create a MetaData container object that binds features of the database being described
  md_obj = MetaData()

  #Define the table "user"
  user = Table('user',md_obj,
    Column('user_id', Integer, primary_key=True),
    Column('user_name',String(30),nullable=False)
  )

  #Define the table "tweets"
  tweets = Table('tweets',md_obj,
    Column('tweet_id', Integer, primary_key=True),
    Column('user_id', Integer, ForeignKey("user.user_id"), nullable=False),
    Column('tweet', String(280), nullable=False)
  )

  #Define the table "profanity_index"
  profanity_index = Table('profanity_index',md_obj,
    Column('tweet_id', Integer, ForeignKey("tweets.tweet_id"), primary_key=True),
    Column("profanity", Float)
  )

  #Connect the engine and create a connection object
  conn = db.connect()

  #Return the connection object and the table objects
  return conn,user,tweets,profanity_index