# Licensed under a 3-clause BSD style license - see LICENSE.rst
import numpy as np
import astropy.units as u
from gammapy.maps import Map, MapCoord, WcsGeom, MapAxes
from gammapy.modeling.models import PowerLawSpectralModel
from gammapy.utils.random import InverseCDFSampler, get_random_state
from gammapy.utils.gauss import Gauss2DPDF
from .kernel import PSFKernel
from .table import EnergyDependentTablePSF
from ..core import IRFMap

__all__ = ["PSFMap"]

class PSFMap(IRFMap):
    """Class containing the Map of PSFs and allowing to interact with it.

    Parameters
    ----------
    psf_map : `~gammapy.maps.Map`
        the input PSF Map. Should be a Map with 2 non spatial axes.
        rad and true energy axes should be given in this specific order.
    exposure_map : `~gammapy.maps.Map`
        Associated exposure map. Needs to have a consistent map geometry.

    Examples
    --------
    ::

        import numpy as np
        from astropy import units as u
        from astropy.coordinates import SkyCoord
        from gammapy.maps import Map, WcsGeom, MapAxis
        from gammapy.irf import EnergyDependentMultiGaussPSF, EffectiveAreaTable2D
        from gammapy.makers.utils import make_psf_map, PSFMap, make_map_exposure_true_energy

        # Define energy axis. Note that the name is fixed.
        energy_axis = MapAxis.from_edges(np.logspace(-1., 1., 4), unit='TeV', name='energy')
        # Define rad axis. Again note the axis name
        rads = np.linspace(0., 0.5, 100) * u.deg
        rad_axis = MapAxis.from_edges(rads, unit='deg', name='theta')

        # Define parameters
        pointing = SkyCoord(0., 0., unit='deg')
        max_offset = 4 * u.deg

        # Create WcsGeom
        geom = WcsGeom.create(binsz=0.25*u.deg, width=10*u.deg, skydir=pointing, axes=[rad_axis, energy_axis])

        # Extract EnergyDependentTablePSF from CTA 1DC IRF
        filename = '$GAMMAPY_DATA/cta-1dc/caldb/data/cta/1dc/bcf/South_z20_50h/irf_file.fits'
        psf = EnergyDependentMultiGaussPSF.read(filename, hdu='POINT SPREAD FUNCTION')
        psf3d = psf.to_psf3d(rads)
        aeff2d = EffectiveAreaTable2D.read(filename, hdu='EFFECTIVE AREA')

        # Create the exposure map
        exposure_geom = geom.to_image().to_cube([energy_axis])
        exposure_map = make_map_exposure_true_energy(pointing, "1 h", aeff2d, exposure_geom)

        # create the PSFMap for the specified pointing
        psf_map = make_psf_map(psf3d, pointing, geom, max_offset, exposure_map)

        # Get an EnergyDependentTablePSF at any position in the image
        psf_table = psf_map.get_energy_dependent_table_psf(SkyCoord(2., 2.5, unit='deg'))

        # Write map to disk
        psf_map.write('psf_map.fits')
    """
    tag = "psf_map"
    required_axes = ["rad", "energy_true"]

    def __init__(self, psf_map, exposure_map=None):
        super().__init__(irf_map=psf_map, exposure_map=exposure_map)

    @property
    def psf_map(self):
        return self._irf_map

    @psf_map.setter
    def psf_map(self, value):
        self._irf_map = value

    def get_energy_dependent_table_psf(self, position=None):
        """Get energy-dependent PSF at a given position.

        By default the PSF at the center of the map is returned.

        Parameters
        ----------
        position : `~astropy.coordinates.SkyCoord`
            the target position. Should be a single coordinates

        Returns
        -------
        psf_table : `~gammapy.irf.EnergyDependentTablePSF`
            the table PSF
        """
        if position is None:
            position = self.psf_map.geom.center_skydir

        if position.size != 1:
            raise ValueError(
                "EnergyDependentTablePSF can be extracted at one single position only."
            )

        energy = self.psf_map.geom.axes["energy_true"].center
        rad = self.psf_map.geom.axes["rad"].center

        coords = {
            "skycoord": position,
            "energy_true": energy.reshape((-1, 1, 1, 1)),
            "rad": rad.reshape((1, -1, 1, 1)),
        }

        data = self.psf_map.interp_by_coord(coords)
        psf_values = u.Quantity(data[:, :, 0, 0], unit=self.psf_map.unit, copy=False)

        if self.exposure_map is not None:
            coords = {
                "skycoord": position,
                "energy_true": energy.reshape((-1, 1, 1)),
                "rad": 0 * u.deg,
            }
            data = self.exposure_map.interp_by_coord(coords).squeeze()
            exposure = data * self.exposure_map.unit
        else:
            exposure = None

        # Beware. Need to revert rad and energies to follow the TablePSF scheme.
        return EnergyDependentTablePSF(
            axes=self.psf_map.geom.axes[["energy_true", "rad"]],
            data=psf_values.value,
            unit=psf_values.unit,
            exposure=exposure,
        )

    def get_psf_kernel(self, position, geom, max_radius=None, factor=4):
        """Returns a PSF kernel at the given position.

        The PSF is returned in the form a WcsNDMap defined by the input Geom.

        Parameters
        ----------
        position : `~astropy.coordinates.SkyCoord`
            the target position. Should be a single coordinate
        geom : `~gammapy.maps.Geom`
            the target geometry to use
        max_radius : `~astropy.coordinates.Angle`
            maximum angular size of the kernel map
        factor : int
            oversampling factor to compute the PSF

        Returns
        -------
        kernel : `~gammapy.irf.PSFKernel`
            the resulting kernel
        """
        if position is None:
            position = self.psf_map.geom.center_skydir

        table_psf = self.get_energy_dependent_table_psf(position)

        if max_radius is None:
            max_radius = np.max(table_psf.axes["rad"].center)
            min_radius_geom = np.min(geom.width) / 2.0
            max_radius = min(max_radius, min_radius_geom)

        return PSFKernel.from_table_psf(table_psf, geom, max_radius, factor)

    def containment_radius_map(self, energy, fraction=0.68):
        """Containment radius map.

        Parameters
        ----------
        energy : `~astropy.units.Quantity`
            Scalar energy at which to compute the containment radius
        fraction : float
            the containment fraction (range: 0 to 1)

        Returns
        -------
        containment_radius_map : `~gammapy.maps.Map`
            Containment radius map
        """
        coords = self.psf_map.geom.to_image().get_coord().skycoord.flatten()
        m = Map.from_geom(self.psf_map.geom.to_image(), unit="deg")

        for coord in coords:
            psf_table = self.get_energy_dependent_table_psf(coord)
            containment_radius = psf_table.containment_radius(
                energy_true=energy, fraction=fraction
            )
            m.fill_by_coord(coord, np.atleast_1d(containment_radius))

        return m

    @classmethod
    def from_geom(cls, geom):
        """Create psf map from geom.

        Parameters
        ----------
        geom : `Geom`
            PSF map geometry.

        Returns
        -------
        psf_map : `PSFMap`
            Point spread function map.
        """
        geom_exposure_psf = geom.squash(axis_name="rad")
        exposure_psf = Map.from_geom(geom_exposure_psf, unit="m2 s")
        psf_map = Map.from_geom(geom, unit="sr-1")
        return cls(psf_map, exposure_psf)

    def sample_coord(self, map_coord, random_state=0):
        """Apply PSF corrections on the coordinates of a set of simulated events.

        Parameters
        ----------
        map_coord : `~gammapy.maps.MapCoord` object.
            Sequence of coordinates and energies of sampled events.
        random_state : {int, 'random-seed', 'global-rng', `~numpy.random.RandomState`}
            Defines random number generator initialisation.
            Passed to `~gammapy.utils.random.get_random_state`.

        Returns
        -------
        corr_coord : `~gammapy.maps.MapCoord` object.
            Sequence of PSF-corrected coordinates of the input map_coord map.
        """

        random_state = get_random_state(random_state)
        rad_axis = self.psf_map.geom.axes["rad"]

        coord = {
            "skycoord": map_coord.skycoord.reshape(-1, 1),
            "energy_true": map_coord["energy_true"].reshape(-1, 1),
            "rad": rad_axis.center,
        }

        pdf = (
            self.psf_map.interp_by_coord(coord)
            * rad_axis.center.value
            * rad_axis.bin_width.value
        )

        sample_pdf = InverseCDFSampler(pdf, axis=1, random_state=random_state)
        pix_coord = sample_pdf.sample_axis()
        separation = rad_axis.pix_to_coord(pix_coord)

        position_angle = random_state.uniform(360, size=len(map_coord.lon)) * u.deg

        event_positions = map_coord.skycoord.directional_offset_by(
            position_angle=position_angle, separation=separation
        )
        return MapCoord.create(
            {"skycoord": event_positions, "energy_true": map_coord["energy_true"]}
        )

    @classmethod
    def from_energy_dependent_table_psf(cls, table_psf, geom=None):
        """Create PSF map from table PSF object.

        Helper function to create an allsky PSF map from
        table PSF, which does not depend on position.

        Parameters
        ----------
        table_psf : `EnergyDependentTablePSF`
            Table PSF

        Returns
        -------
        psf_map : `PSFMap`
            Point spread function map.
        """
        if geom is None:
            geom = WcsGeom.create(
                npix=(2, 1),
                proj="CAR",
                binsz=180,
                axes=table_psf.axes.reverse,
            )
        coords = geom.get_coord()

        data = table_psf.evaluate(
            energy_true=coords["energy_true"], rad=coords["rad"]
        )
        psf_map = Map.from_geom(geom, data=data.to_value("sr-1"), unit="sr-1")

        geom_exposure = geom.squash(axis_name="rad")

        exposure_map = Map.from_geom(geom_exposure, unit="cm2 s", data=1)
        data = table_psf.exposure
        exposure_map.quantity = exposure_map.data * data.reshape((-1, 1, 1, 1))
        return cls(psf_map=psf_map, exposure_map=exposure_map)

    @classmethod
    def from_gauss(cls, energy_axis_true, rad_axis=None, sigma=0.1 * u.deg):
        """Create all -sky PSF map from Gaussian width.

        This is used for testing and examples.

        The width can be the same for all energies
        or be an array with one value per energy node.
        It does not depend on position.

        Parameters
        ----------
        energy_axis_true : `~gammapy.maps.MapAxis`
            True energy axis.
        rad_axis : `~gammapy.maps.MapAxis`
            Offset angle wrt source position axis.
        sigma : `~astropy.coordinates.Angle`
            Gaussian width.

        Returns
        -------
        psf_map : `PSFMap`
            Point spread function map.
        """
        from gammapy.datasets.map import RAD_AXIS_DEFAULT

        if rad_axis is None:
            rad_axis = RAD_AXIS_DEFAULT.copy()

        axes = MapAxes([energy_axis_true, rad_axis])
        coords = axes.get_coord()

        sigma = np.broadcast_to(u.Quantity(sigma), energy_axis_true.nbin, subok=True)
        gauss = Gauss2DPDF(sigma=sigma.reshape((-1, 1)))
        data = gauss(coords["rad"])

        table_psf = EnergyDependentTablePSF(
            axes=axes, unit=data.unit, data=data.value
        )
        return cls.from_energy_dependent_table_psf(table_psf)

    def to_image(self, spectrum=None, keepdims=True):
        """Reduce to a 2-D map after weighing
        with the associated exposure and a spectrum

        Parameters
        ----------
        spectrum : `~gammapy.modeling.models.SpectralModel`, optional
            Spectral model to compute the weights.
            Default is power-law with spectral index of 2.
        keepdims : bool, optional
            If True, the energy axis is kept with one bin.
            If False, the axis is removed

        Returns
        -------
        psf_out : `PSFMap`
            `PSFMap` with the energy axis summed over
        """
        from gammapy.makers.utils import _map_spectrum_weight

        if spectrum is None:
            spectrum = PowerLawSpectralModel(index=2.0)

        exp_weighed = _map_spectrum_weight(self.exposure_map, spectrum)
        exposure = exp_weighed.sum_over_axes(
            axes_names=["energy_true"], keepdims=keepdims
        )

        psf_data = exp_weighed.data * self.psf_map.data / exposure.data
        psf_map = Map.from_geom(geom=self.psf_map.geom, data=psf_data, unit="sr-1")

        psf = psf_map.sum_over_axes(axes_names=["energy_true"], keepdims=keepdims)
        return self.__class__(psf_map=psf, exposure_map=exposure)

import numpy,import astropy,from astropy


import numpy as np
import random
from ForgER.sum_tree import SumSegmentTree, MinSegmentTree

class PrioritizedReplayBuffer:
    def __init__(self, size, env_dict, alpha=0.6, eps=1e-6):
        """Create Prioritized Replay buffer.
        Parameters
        ----------
        size: int
            Max number of transitions to store in the buffer. When the buffer
            overflows the old memories are dropped.
        alpha: float
            how much prioritization is used
            (0 - no prioritization, 1 - full prioritization)
        See Also
        --------
        ReplayBuffer.__init__
        """
        #super(PrioritizedReplayBuffer, self).__init__(size, env_dict)
        self._storage = []
        self._maxsize = size
        self._next_idx = 0
        self.env_dict = env_dict

        assert alpha >= 0
        self._alpha = alpha
        self._eps = eps

        it_capacity = 1
        while it_capacity < size:
            it_capacity *= 2

        self._it_sum = SumSegmentTree(it_capacity)
        self._it_min = MinSegmentTree(it_capacity)
        self._max_priority = 1.0

    def get_stored_size(self):
        return len(self._storage)

    def get_buffer_size(self):
        return self._maxsize

    def add(self, kwargs):
        #print("kwargs: " + str(kwargs))

        """See ReplayBuffer.store_effect"""
        idx = self._next_idx
        #super().add(kwargs=kwargs)
        data = kwargs
        if self._next_idx >= len(self._storage):
            self._storage.append(data)
        else:
            self._storage[self._next_idx] = data

        self._next_idx = (self._next_idx + 1) % self._maxsize

        self._it_sum[idx] = self._max_priority ** self._alpha
        self._it_min[idx] = self._max_priority ** self._alpha

    def _sample_proportional(self, batch_size):
        res = []
        p_total = self._it_sum.sum(0, len(self._storage) - 1)
        every_range_len = p_total / batch_size
        for i in range(batch_size):
            mass = random.random() * every_range_len + i * every_range_len
            idx = self._it_sum.find_prefixsum_idx(mass)
            res.append(idx)

        return res

    def sample(self, batch_size, beta=0.4):
        """Sample a batch of experiences.
        compared to ReplayBuffer.sample
        it also returns importance weights and idxes
        of sampled experiences.
        Parameters
        ----------
        batch_size: int
            How many transitions to sample.
        beta: float
            Priority level
        Returns
        -------
        encoded_sample: dict of np.array
            Array of shape(batch_size, ...) and dtype np.*32
        weights: np.array
            Array of shape (batch_size,) and dtype np.float32
            denoting importance weight of each sampled transition
        idxes: np.array
            Array of shape (batch_size,) and dtype np.int32
            idexes in buffer of sampled experiences
        """
        idxes = self._sample_proportional(batch_size)
        it_sum = self._it_sum.sum()
        it_min = self._it_min.min()
        p_min = it_min / it_sum
        max_weight = (p_min * len(self._storage)) ** (- beta)
        p_sample = np.array([self._it_sum[idx] / it_sum for idx in idxes])
        weights = (p_sample*len(self._storage)) ** (- beta)
        weights = weights / max_weight
        encoded_sample = self._encode_sample(idxes)
        encoded_sample['weights'] = weights
        encoded_sample['indexes'] = idxes

        return encoded_sample

    @property
    def first_transition(self):
        return self._storage[0]

    def _encode_sample(self, idxes):
        batch = {key: list() for key in self.first_transition.keys()}
        #print("batch: " + str(batch))
        for i in idxes:
            data = self._storage[i]
            for key, value in data.items():
                batch[key].append(np.array(value))

        for key, value in batch.items():
            #print("type(key): " + str(type(key)))
            #print("self.env_dict[key]: " + str(self.env_dict[key]))
            #print("len(value): " + str(len(value)))
            #print("")
            #test = self.env_dict[key]
            test = np.array(value, dtype=self.env_dict[key])
            batch[key] = np.array(value, dtype=self.env_dict[key])

        return batch

    def update_priorities(self, idxes, priorities):
        """Update priorities of sampled transitions.
        sets priority of transition at index idxes[i] in buffer
        to priorities[i].
        Parameters
        ----------
        idxes: [int]
            List of idxes of sampled transitions
        priorities: [float]
            List of updated priorities corresponding to
            transitions at the sampled idxes denoted by
            variable `idxes`.
        """
        assert len(idxes) == len(priorities)
        for idx, priority in zip(idxes, priorities):
            assert priority >= 0
            assert 0 <= idx < len(self._storage)
            self._it_sum[idx] = (priority + self._eps) ** self._alpha
            self._it_min[idx] = (priority + self._eps) ** self._alpha

            self._max_priority = max(self._max_priority, priority)

class AggregatedBuff:
    def __init__(self, size, env_dict, episodes_to_decay=50, min_demo_proportion=0.0):
        buffer_base = PrioritizedReplayBuffer

        self._maxsize = size

        #print("size: " + str(size))
        #print("env_dict: " + str(env_dict))
        '''
        env_dict: {'n_reward': {'dtype': 'float32'}, 'to_demo': {'dtype': 'float32'}, 'demo': {'dtype': 'float32'}, 'n_done': {'dtype': 'bool'},
        'actual_n': {'dtype': 'float32'}, 'infrared': {'dtype': dtype('int64'), 'shape': ()}, 'next_infrared': {'dtype': dtype('int64'), 'shape': ()},
        'next_camera': {'dtype': dtype('float32'), 'shape': (512, 512, 3)}, 'done': {'dtype': 'bool'}, 'next_state': {'dtype': 'float32'},
        'n_lidar': {'dtype': dtype('float32'), 'shape': (1,)}, 'n_camera': {'dtype': dtype('float32'), 'shape': (512, 512, 3)}, 'next_lidar': {'dtype': dtype('float32'),
        'shape': (1,)}, 'lidar': {'dtype': dtype('float32'), 'shape': (1,)}, 'n_infrared': {'dtype': dtype('int64'), 'shape': ()}, 'state': {'dtype': 'float32'},
        'camera': {'dtype': dtype('float32'), 'shape': (512, 512, 3)}, 'n_state': {'dtype': 'float32'}, 'action': {'dtype': 'int32', 'shape': 1},
        'reward': {'dtype': 'float32'}}
        '''

        self.demo_kwargs = {'size': size, 'env_dict': env_dict, 'eps': 1.0}
        self.demo_buff = buffer_base(size=size, env_dict=env_dict, eps=1.0)
        self.replay_buff = buffer_base(size=size, env_dict=env_dict, eps=1e-3)
        self.episodes_to_decay = episodes_to_decay
        self.episodes_done = 0
        self.min_demo_proportion = min_demo_proportion

    def add(self, kwargs, to_demo=0):
        #print("kwargs.keys(): " + str(kwargs.keys()))

        if to_demo:
            self.demo_buff.add(kwargs)
        else:
            #print("kwargs: " + str(kwargs))
            self.replay_buff.add(kwargs)
            if kwargs['done']:
                self.episodes_done += 1
                if self.demo_buff.get_stored_size() > 0\
                        and self.episodes_done > self.episodes_to_decay\
                        and self.min_demo_proportion == 0.:
                    self.demo_buff = PrioritizedReplayBuffer(self.demo_kwargs)

    def free_demo(self):
        self.demo_buff = PrioritizedReplayBuffer(self.demo_kwargs)

    @property
    def proportion(self):
        if self.episodes_to_decay == 0:
            proportion = 1. - self.min_demo_proportion
        else:
            proportion = min(1. - self.min_demo_proportion, self.episodes_done / self.episodes_to_decay)

        proportion = max(proportion, float(self.demo_buff.get_stored_size() == 0))

        return proportion

    def sample(self, n=32, beta=0.4):
        #print("n: " + str(n))
        agent_n = int(n * self.proportion)
        demo_n = n - agent_n

        if demo_n > 0 and agent_n > 0:
            demo_samples = self.demo_buff.sample(demo_n, beta)
            replay_samples = self.replay_buff.sample(agent_n, beta)
            samples = {key: np.concatenate((replay_samples[key], demo_samples[key]))
                       for key in replay_samples.keys()}
        elif agent_n == 0:
            samples = self.demo_buff.sample(demo_n, beta)
        else:
            samples = self.replay_buff.sample(agent_n, beta)

        samples = {key: np.squeeze(value) for key, value in samples.items()}

        return samples

    def update_priorities(self, indexes, priorities):
        n = len(indexes)
        agent_n = int(n * self.proportion)
        demo_n = n - agent_n
        if demo_n != 0:
            self.demo_buff.update_priorities(indexes[agent_n:], priorities[agent_n:])

        if agent_n != 0:
            self.replay_buff.update_priorities(indexes[:agent_n], priorities[:agent_n])

    def get_stored_size(self):
        return self.replay_buff.get_stored_size()

    def get_buffer_size(self):
        return self._maxsize


import numpy as np
import os
import matplotlib.pyplot as plt
import librosa
#import librosa.display
import madmom
from pydub import AudioSegment
import math
import json

NEED_COPY_FOLDER = False
filename='_original.mp3'
pianofile='piano.mp3'
in_dir = 'D:/2019/Summer/intern/audio_word/pop2jazz/'
in_dir += 'transcription_with_disjoint_notes.from_separated.soft_jazz_trio/20190403_095002.20190524_082516'
cut_multitrack = True
folder_dict = {}

'''
def copy_folder_structure(inputpath, outputpath):
    for dirpath, dirnames, filenames in os.walk(inputpath):
        structure = os.path.join(outputpath, os.path.relpath(dirpath, inputpath))
        if not os.path.isdir(structure): os.mkdir(structure)
        else: print("Folder does already exits!")
'''
def find_downbeats(source):
    proc = madmom.features.downbeats.DBNDownBeatTrackingProcessor(beats_per_bar=[4, 4], fps=100)
    act = madmom.features.downbeats.RNNDownBeatProcessor()(source)
    a = proc(act)
    return a

def load_audio(source):
    sig, sr = librosa.core.load(source) # sr = 22050
    #sig = AudioSegment.from_file(source)
    #sound2 = AudioSegment.from_file(piano_file)
    #print('Time:{}, Length:{}'.format(len(sig)/sr, len(sig)))
    return sig, sr
    #return sig

def segment(source, pianotrack, outputpath, outputpath_p, index):#, sr=44100):
    # Pydub do things in sr=1000
    a = find_downbeats(source)
    print('Track downbeats done ...')
    #sig = load_audio(source)

    sig_p, sr_p = load_audio(pianotrack)
    if(cut_multitrack ==True):
        sig, sr = load_audio(source)
    print('Load audio done ...')
    #sigs = []
    start = 0
    #has_record=False
    for i, c in enumerate(a):
        if c[-1] == 1:

            #if (has_record==False):
            #    has_record=True
            #    start=int(c[0]*sr)
            #    continue
            output_name_p = os.path.join(outputpath_p,'cut{:03d}_{:03d}.wav'.format(index, int(i/4)+1))
            #print(output_name_p)
            #sigs.append(sig[start : int(i[0]*sr)])
            cut_p = sig_p[start : int(c[0]*sr_p)]
            #cut.export(output_name,format='wav')
            librosa.output.write_wav(output_name_p, cut_p, sr_p)

            if(cut_multitrack==True):
                output_name = os.path.join(outputpath,'cut{:03d}_{:03d}.wav'.format(index, int(i/4)+1))
                cut = sig[start : int(c[0]*sr)]
                librosa.output.write_wav(output_name, cut, sr)
            start = int(c[0]*sr_p)
            #librosa.output.write_wav(output_name, sigs[i], sr)
    #sigs.append(sig[start:])
    #print(len(sigs)) # num of bars

def dictToTxt(in_dict, out_txt):
    with open(out_txt, 'w') as file:
        file.write(json.dumps(in_dict))

def txtToDict(in_txt):
    with open(in_txt, 'r') as file:
        a = json.load(file)
    return a
###########################################
out_dir_piano='./music_segments_piano'
out_dir_all='./music_segments_all_track'
# if NEED_COPY_FOLDER:
#     copy_folder_structure(in_dir, out_dir)
#     NEED_COPY_FOLDER = False

index = 0
for f in os.listdir(in_dir):
    if os.path.isdir(os.path.join(in_dir,f)):
        print(f) #folder
        for ff in os.listdir(os.path.join(in_dir,f)):
            print('---',ff) #sub-folder
            fder = os.path.join(os.path.join(in_dir,f),ff)
            print(fder)
            if os.path.isdir(fder):
                folder_dict['{:03d}'.format(index)] = os.path.join(f+'/',ff)
                print(folder_dict)
                save_dir_piano = os.path.join(out_dir_piano, '{:03d}'.format(index))
                save_dir_all = os.path.join(out_dir_all, '{:03d}'.format(index))
                print(save_dir_piano)

                if not os.path.isdir(save_dir_piano):
                    os.mkdir(save_dir_piano)
                if not os.path.isdir(save_dir_all):
                    os.mkdir(save_dir_all)
                segment(os.path.join(fder, filename),os.path.join(fder, pianofile), save_dir_all, save_dir_piano, index)
                print('done')
                index += 1
                # for file in os.listdir(fder):
                #     print(file)

    dictToTxt(folder_dict, 'file.txt')

#d = txtToDict('file.txt')
#print(',,,',d)


import os
import gc
import csv
import torch
import argparse
import numpy as np
import torch.nn.functional as F
from torch.autograd import Variable
from sklearn.metrics import f1_score
import torch.backends.cudnn as cudnn
from torch.utils.data import DataLoader

from models import *
from dataset import *
from utils import progress_bar

parser = argparse.ArgumentParser(description='PyTorch OffenseEval - run dataset.py first for word embeddings')
parser.add_argument('--lr', default=0.01, type=float, help='learning rate') # NOTE :  change for diff models
parser.add_argument('--batch_size', default=50, type=int)
parser.add_argument('--resume', '-r', type=int, default=0, help='resume from checkpoint')
parser.add_argument('--epochs', '-e', type=int, default=50, help='Number of epochs to train.')
parser.add_argument('--subtask', default='C', help="Sub-task for OffensEval")
parser.add_argument('--embedding_length', default=50, type=int)

args = parser.parse_args()

device = 'cuda' if torch.cuda.is_available() else 'cpu'
best_acc, tsepoch, tstep = 0., 0, 0#[0., 0., 0.], 0, 0

criterion = torch.nn.CrossEntropyLoss(reduction='none')

print('==> Preparing data..')

'''def collate_fn(data):
    data = list(filter(lambda x: type(x[1]) != int, data))
    audios, captions = zip(*data)
    data = None
    del data
    audios = torch.stack(audios, 0)
    return audios, captions'''

classes = {"A" : 3, "B" : 3, "C" : 4}
print('==> Creating network..')
net = AttentionModel(args.batch_size, [3, 3, 4], 25, args.embedding_length)
net = net.to(device)

if(args.resume):
    if(os.path.isfile('../save/network.ckpt')):
        net.load_state_dict(torch.load('../save/network.ckpt'))
        print('==> Network : loaded')

    if(os.path.isfile("../save/info.txt")):
        with open("../save/info.txt", "r") as f:
            tsepoch, tstep = (int(i) for i in str(f.read()).split(" "))
        print("=> Network : prev epoch found")
else :
    with open("../save/logs/train_loss.log", "w+") as f:
        pass

def train_network(epoch):
    global tstep
    global best_acc
    print('\n=> Epoch: {}'.format(epoch))
    net.train()

    dataset = OffenseEval(path='/home/nevronas/Projects/Personal-Projects/Dhruv/OffensEval/dataset/train-v1/offenseval-training-v1.tsv')
    dataloader = DataLoader(dataset, batch_size=args.batch_size, shuffle=True) #, collate_fn=collate_fn)
    dataloader = iter(dataloader)

    train_loss, accu1, accu2, accu3 = 0.0, 0.0, 0.0, 0.0
    le = len(dataloader) - 1
    params = net.parameters()
    optimizer = torch.optim.Adam(params, lr=args.lr)

    # Training for task A
    for i in range(tstep, le):
        contents = next(dataloader)
        inputs = contents[0].type(torch.FloatTensor).to(device)
        target_a = contents[1].type(torch.LongTensor).to(device) #, target_b, target_c = [contents[i].type(torch.LongTensor).to(device) for i in range(1, len(contents), 1)]
        #targets = [target_a, target_b, target_c]

        suban = target_a.detach().cpu().numpy() #, subbn, subcn = target_a.detach().cpu().numpy(), target_b.detach().cpu().numpy(), target_c.detach().cpu().numpy()
        mask1 = np.where(suban == 0) #, mask2, mask3 = np.where(suban == 0), np.where(subbn == 0), np.where(subcn == 0)

        optimizer.zero_grad()
        y_preds = net(inputs)
        #preds = [torch.max(y_pred, 1)[0].type(torch.LongTensor) for y_pred in y_preds]

        l1o = criterion(y_preds, target_a) #, l2o, l3o = criterion(y_preds[0], target_a), criterion(y_preds[1], target_b), criterion(y_preds[2], target_c)
        l1 = l1o.detach().cpu().numpy() #, l2, l3 = l1o.detach().cpu().numpy(), l2o.detach().cpu().numpy(), l3o.detach().cpu().numpy()
        l1[mask1] = 0 #, l2[mask2], l3[mask3] = 0, 0, 0
        l1 = torch.Tensor(l1).to(device) #, l2, l3 = torch.Tensor(l1).to(device), torch.Tensor(l2).to(device), torch.Tensor(l3).to(device)

        loss = torch.mean((l1 * l1o) / 2) #+ torch.mean((l2 * l2o) / 2) + torch.mean((l3 * l3o) / 2)
        tl = loss.item()
        loss.backward()
        optimizer.step()

        acc1 = f1_score(target_a.detach().cpu().numpy(), torch.max(y_preds, 1)[0].type(torch.LongTensor).detach().cpu().numpy(), average='macro')
        #acc2 = f1_score(target_b.detach().cpu().numpy(), preds[1].detach().cpu().numpy(), average='macro')
        #acc3 = f1_score(target_c.detach().cpu().numpy(), preds[2].detach().cpu().numpy(), average='macro')

        train_loss += tl
        accu1 += acc1
        #accu2 += acc2
        #accu3 += acc3

        gc.collect()
        torch.cuda.empty_cache()

        torch.save(net.state_dict(), '../save/network.ckpt')
        with open("../save/info.txt", "w+") as f:
            f.write("{} {}".format(epoch, i))

        with open("../save/logs/train_loss.log", "a+") as lfile:
            lfile.write("{}\n".format(tl))

        progress_bar(i, len(dataloader), 'Loss: {}, F1s: {} '.format(tl, acc1))#, acc2, acc3)) # "{} - {}"

    tstep = 0
    del dataloader
    print('=> Network : Epoch [{}/{}], Loss:{:.4f}, F1:{:.4f}'.format(epoch + 1, args.epochs, train_loss / le, accu1 / le)) #,  accu2 / le, accu3 / le))
    #accu = [accu1/le, accu2/le, accu3/le]
    #best_acc = [max(best_acc[i], accu[i]) for i in range(3)]
    old_best = best_acc
    best_acc = max(best_acc, accu1/le)
    if(best_acc != old_best):
        torch.save(net.state_dict(), '../save/best.ckpt')
    print("Best Metrics : {}".format(best_acc))

def test():
    global net
    net.load_state_dict(torch.load('../save/success/1.ckpt'))

    dataset = OffenseEval(path='/home/nevronas/Projects/Personal-Projects/Dhruv/OffensEval/dataset/testset-taska.tsv', train=False)
    dataloader = DataLoader(dataset, batch_size=args.batch_size) #, collate_fn=collate_fn)
    dataloader = iter(dataloader)
    test_dict = ['NOT', 'OFF']

    with open('../save/test.tsv', 'w+') as f:
        tsv_writer = csv.writer(f, delimiter='\t')
        tsv_writer.writerow(['id', 'tweet', 'subtask_a'])

        for i in range(0, len(dataloader) - 1):
            contents = next(dataloader)
            inputs = contents[2].type(torch.FloatTensor).to(device)
            y_preds = net(inputs)
            clas = torch.max(y_preds, 1)[0].type(torch.LongTensor) - 1
            clas = clas.tolist()

            for i in range(len(clas)):
                if(int(clas[i]) > 1):
                    clas[i] = 1
                print(test_dict[int(clas[i])])
                tsv_writer.writerow([contents[0][i], contents[1][i], test_dict[int(clas[i])] ])

#for epoch in range(tsepoch, tsepoch + args.epochs):
#    train_network(epoch)

test()


import argparse
import csv
import math
import os
import pathlib
import subprocess

import torch
import shutil
import numpy as np
from PIL import Image
from tqdm import tqdm

import aegan.aegan as aegan

def revert(im):
    im -= im.min()
    im /= im.max()
    im = im.numpy().transpose((1, 2, 0))
    im = np.array(im * 255, dtype=np.uint8)
    return im

def main(args):
    noise_fn = lambda x: torch.randn((x, 16), device='cpu')

    with open("fid.csv", ("wa" if args.append else "w")) as f:
        csv_writer = csv.DictWriter(f, fieldnames=['iternum', 'FID'])
        if not args.append:
            csv_writer.writeheader()

        paths = sorted(os.listdir(args.checkpoints_dir))
        paths = paths[args.from_n:]
        for path in tqdm(paths):
            print(path)
            path = pathlib.Path(path)
            gan = aegan.AEGAN(16, noise_fn, None, device="cuda")

            gan.generator.load_state_dict(torch.load(args.checkpoints_dir / path / "generator.pt", map_location="cpu"))
            gan.generator.train(False)
            generator = gan.generator

            latent_vec = noise_fn(args.num).cuda()
            batch_size = 50
            samples = []
            with torch.no_grad():
                for i in range(0, math.ceil(args.num / batch_size)):
                    l = i * batch_size
                    r = min(args.num, i * batch_size + batch_size)
                    samples.append(generator(latent_vec[l:r]))
            samples = torch.cat(samples)
            samples = samples.cpu()

            data_path = args.checkpoints_dir / path / "gen_dataset"
            if data_path.exists():
                shutil.rmtree(data_path, ignore_errors=True)
            os.makedirs(data_path, exist_ok=False)

            for i in range(args.num):
                im = Image.fromarray(revert(samples[i]))
                im.save(args.checkpoints_dir / path / "gen_dataset" / f"{i:04d}.png")

            result = subprocess.run(["pytorch-fid",
                                     "--batch-size", "16",
                                     "--device", "cuda",
                                     args.original_data, data_path], stdout=subprocess.PIPE)
            fid = result.stdout.decode().split()[-1]
            csv_writer.writerow({"iternum": int(str(path)[-5:]), "FID": fid})
            f.flush()

if __name__ == '__main__':
    parser = argparse.ArgumentParser()

    parser.add_argument("checkpoints_dir", type=pathlib.Path)
    parser.add_argument("original_data", type=pathlib.Path)
    parser.add_argument("-n", "--num", type=int, required=True)
    parser.add_argument("--append", action="store_true")
    parser.add_argument("--from_n", type=int, required=True)

    main(parser.parse_args())


import argparse
import cv2
import numpy as np
import torch
import mmcv
from mmcv.runner import load_checkpoint, parallel_test, obj_from_dict
from mmcv.parallel import scatter, collate, MMDataParallel

from mmaction import datasets
from mmaction.datasets import build_dataloader
from mmaction.models import build_recognizer, recognizers
from mmaction.core.evaluation.accuracy import (softmax, top_k_accuracy,
                                               mean_class_accuracy)

def inference(model, images):
    result = model([1], None, return_loss=False, img_group_0=torch.Tensor(images))
    return result

def single_test(model, data_loader):
    model.eval()
    results = []
    dataset = data_loader.dataset
    # prog_bar = mmcv.ProgressBar(len(dataset))
    for i, data in enumerate(data_loader):
        with torch.no_grad():
            # result = model(return_loss=False, **data)
            result = inference(model, data['img_group_0'])
        results.append(result)

        # batch_size = data['img_group_0'].data[0].size(0)
        # for _ in range(batch_size):
        #     prog_bar.update()
    return results

def _data_func(data, device_id):
    data = scatter(collate([data], samples_per_gpu=1), [device_id])[0]
    return dict(return_loss=False, rescale=True, **data)

def parse_args():
    parser = argparse.ArgumentParser(description='Test an action recognizer')
    parser.add_argument('video_path', help='path to test video')
    parser.add_argument('config', help='test config file path')
    parser.add_argument('checkpoint', help='checkpoint file')
    parser.add_argument(
        '--gpus', default=1, type=int, help='GPU number used for testing')
    parser.add_argument(
        '--proc_per_gpu',
        default=1,
        type=int,
        help='Number of processes per GPU')
    parser.add_argument('--use_softmax', action='store_true',
                        help='whether to use softmax score')
    args = parser.parse_args()
    return args

def main():
    args = parse_args()

    print(args.video_path)
    cap = cv2.VideoCapture(args.video_path)
    frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    print('Video length: ', frame_count)
    frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    buf = np.empty((frame_count - 1, 224, 224, 3), np.dtype('uint8'))

    fc = 0
    ret = True

    while (fc < frame_count - 1 and ret):
        ret, frame = cap.read()
        buf[fc] = cv2.resize(frame, (224, 224))
        fc += 1
    print('Video loaded')

    cap.release()

    buf = buf.transpose((0, 3, 1, 2))
    buf = np.expand_dims(buf, 0)
    buf = buf.astype(np.float32) - 128

    cfg = mmcv.Config.fromfile(args.config)
    # set cudnn_benchmark
    if cfg.get('cudnn_benchmark', False):
        torch.backends.cudnn.benchmark = True
    cfg.data.test.test_mode = True

    if cfg.data.test.oversample == 'three_crop':
        cfg.model.spatial_temporal_module.spatial_size = 8

    model = build_recognizer(
        cfg.model, train_cfg=None, test_cfg=cfg.test_cfg)
    load_checkpoint(model, args.checkpoint, strict=True)

    model.eval()
    outputs = inference(model, buf)

    pred = np.argmax(outputs, axis=1)[0]
    if hasattr(cfg, 'class_names'):
        print('Class #', pred, ': ', cfg.class_names[pred])
    else:
        print('Class #', pred)

if __name__ == '__main__':
    main()


import tensorflow as tf
physical_devices = tf.config.experimental.list_physical_devices('GPU')
if len(physical_devices) > 0:
    tf.config.experimental.set_memory_growth(physical_devices[0], True)
from absl import app, flags, logging
from absl.flags import FLAGS
import core.utils as utils
from core.yolov4 import filter_boxes
from tensorflow.python.saved_model import tag_constants
import cv2
import numpy as np
from tensorflow.compat.v1 import ConfigProto
from tensorflow.compat.v1 import InteractiveSession
from math import ceil
from time import time

def main(_argv):
    config = ConfigProto()
    config.gpu_options.allow_growth = True
    session = InteractiveSession(config=config)
    FLAGS = {}
    FLAGS["tiny"] = True
    FLAGS["model"] = "yolov4"
    STRIDES, ANCHORS, NUM_CLASS, XYSCALE = utils.load_config(FLAGS)

    score = .25
    iou = 0.5
    input_size = 416
    image_path = "./test_images/00000048.png"
    weights_path = "./checkpoints/yolov4tiny-custom_best-416-tflite-fp16.tflite"

    interpreter = tf.lite.Interpreter(model_path=weights_path)
    interpreter.allocate_tensors()
    input_details = interpreter.get_input_details()
    output_details = interpreter.get_output_details()

    im_name = image_path.split("/")[-1]
    # ext = im_name.split(".")[-1]
    start = time()
    original_image = cv2.imread(image_path)
    h, w, _ = original_image.shape
    h_new = ceil(h/input_size) * input_size
    w_new = ceil(w/input_size) * input_size
    # scaled_h = h_new/h
    # scaled_w = w_new/w
    resized_image = cv2.resize(original_image, (w_new, h_new), cv2.INTER_LINEAR)

    col_list = []
    for i in range(h_new//input_size):
        row_list = []
        for j in range(w_new//input_size):
            tiled_image_original = resized_image[i*input_size : (i+1)*input_size, j*input_size : (j+1)*input_size, :]
            tiled_image = tiled_image_original / 255.
            tiled_image = np.asarray(tiled_image).astype(np.float32)

            interpreter.set_tensor(input_details[0]['index'], [tiled_image])
            interpreter.invoke()
            pred = [interpreter.get_tensor(output_details[i]['index']) for i in range(len(output_details))]
            boxes, pred_conf = filter_boxes(pred[0], pred[1], score_threshold=0.25, input_shape=tf.constant([input_size, input_size]))

            boxes, scores, classes, valid_detections = tf.image.combined_non_max_suppression(
                boxes=tf.reshape(boxes, (tf.shape(boxes)[0], -1, 1, 4)),
                scores=tf.reshape(pred_conf, (tf.shape(pred_conf)[0], -1, tf.shape(pred_conf)[-1])),
                max_output_size_per_class=50,
                max_total_size=50,
                iou_threshold=iou,
                score_threshold=score)
            pred_bbox = [boxes.numpy(), scores.numpy(), classes.numpy(), valid_detections.numpy()]
            row_list.append(utils.draw_bbox(tiled_image_original, pred_bbox))
        col_list.append(cv2.hconcat(row_list))
    scaled_image = cv2.vconcat(col_list)
    reconstructed_image = cv2.resize(scaled_image, (w, h), interpolation = cv2.INTER_AREA)
    end = time()
    session.close()
    #cv2.imshow("predictions", reconstructed_image)
    cv2.imwrite("result_tflite_fp16_"+ im_name, reconstructed_image)
    #cv2.waitKey(0)
    cv2.destroyAllWindows()
    print(end-start)

if __name__ == '__main__':
    try:
        app.run(main)
    except SystemExit:
        pass


# from __future__ import print_function, absolute_import, division
import logging
import multiprocessing
import shutil
import types
from distutils.dir_util import copy_tree
from pathlib import Path
import traceback

import batman
import foldedleastsquares
import lightkurve
import numpy as np
import matplotlib
import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec
import yaml
from lcbuilder.constants import CUTOUT_SIZE
from lightkurve import TessLightCurve
from matplotlib.colorbar import Colorbar
from matplotlib import patches
from astropy.visualization.mpl_normalize import ImageNormalize
from astropy.table import Table
from astropy.io import ascii
import astropy.visualization as stretching
from argparse import ArgumentParser

from scipy import stats

from sherlockpipe import tpfplotter
import six
import sys
import sherlockpipe.LATTE
sys.modules['astropy.extern.six'] = six
sys.modules['LATTE'] = sherlockpipe.LATTE

matplotlib.use('Agg')
import pandas as pd
import os
from os.path import exists
import ast
import csv
from sherlockpipe.LATTE import LATTEutils, LATTEbrew
from os import path
from math import ceil

# get the system path
syspath = str(os.path.abspath(LATTEutils.__file__))[0:-14]
# ---------

# --- IMPORTANT TO SET THIS ----
out = 'pipeline'  # or TALK or 'pipeline'
ttran = 0.1
resources_dir = path.join(path.dirname(__file__))

class Vetter:
    """
    Provides transiting candidate vetting information like centroids and spaceship motion, momentum dumps, neighbours
    curves inspection and more to give a deeper insight on the quality of the candidate signal.
    """
    def __init__(self, object_dir, validate):
        self.args = types.SimpleNamespace()
        self.args.noshow = True
        self.args.north = False
        self.args.o = True
        self.args.mpi = False
        self.args.auto = True
        self.args.save = True
        self.args.nickname = ""  # TODO do we set the sherlock id?
        self.args.FFI = False  # TODO get this from input
        self.args.targetlist = "best_signal_latte_input.csv"
        self.args.new_path = ""  # TODO check what to do with this
        self.object_dir = os.getcwd() if object_dir is None else object_dir
        self.latte_dir = str(Path.home()) + "/.sherlockpipe/latte/"
        if not os.path.exists(self.latte_dir):
            os.mkdir(self.latte_dir)
        self.data_dir = self.object_dir
        self.validation_runs = 5
        self.validate = validate

    def update(self):
        """
        Updates the LATTE metadata to be up to date with the latest TESS information.
        """
        indir = self.latte_dir
        if os.path.exists(indir) and os.path.isdir(indir):
            shutil.rmtree(indir, ignore_errors=True)
        os.makedirs(indir)
        with open("{}/_config.txt".format(indir), 'w') as f:
            f.write(str(indir))
        logging.info("Download the text files required ... ")
        logging.info("Only the manifest text files (~325 M) will be downloaded and no TESS data.")
        logging.info("This step may take a while but luckily it only has to run once... \n")
        if not os.path.exists("{}".format(indir)):
            os.makedirs(indir)
        if not os.path.exists("{}/data".format(indir)):
            os.makedirs("{}/data".format(indir))
        outF = open(indir + "/data/temp.txt", "w")
        outF.write("#all LC file links")
        outF.close()
        outF = open(indir + "/data/temp_tp.txt", "w+")
        outF.write("#all LC file links")
        outF.close()
        LATTEutils.data_files(indir)
        LATTEutils.tp_files(indir)
        LATTEutils.TOI_TCE_files(indir)
        LATTEutils.momentum_dumps_info(indir)

    def __prepare(self, candidate_df):
        """
        Downloads and fills files to be used by LATTE analysis.
        @return: the latte used directory, an open df to be used by LATTE analysis and the tic selected
        """
        # check whether a path already exists
        indir = self.latte_dir
        # SAVE the new output path
        if not os.path.exists("{}/_config.txt".format(indir)):
            self.update()
        candidate_df['id'] = candidate_df['id'].str.replace("TIC ", "")
        TIC_wanted = list(set(candidate_df['id']))
        nlc = len(TIC_wanted)
        logging.info("nlc length: {}".format(nlc))
        logging.info('{}/manifest.csv'.format(str(indir)))
        if exists('{}/manifest.csv'.format(str(indir))):
            logging.info("Existing manifest file found, will skip previously processed LCs and append to end of manifest file")
        else:
            logging.info("Creating new manifest file")
            metadata_header = ['id', 'Marked Transits', 'Sectors', 'RA', 'DEC', 'Solar Rad', 'TMag', 'Teff',
                               'thissector', 'TOI', 'TCE', 'TCE link', 'EB', 'Systematics', 'Background Flux',
                               'Centroids Positions', 'Momentum Dumps', 'Aperture Size', 'In/out Flux', 'Keep',
                               'Comment', 'starttime']
            with open('{}/manifest.csv'.format(str(indir)), 'w') as f:  # save in the photometry folder
                writer = csv.writer(f, delimiter=',')
                writer.writerow(metadata_header)
        return indir, candidate_df, TIC_wanted, candidate_df.iloc[0]["ffi"]

    def __process(self, indir, tic, sectors_in, t0, period, duration, depth, ffi, run, curve):
        """
        Performs the LATTE analysis to generate PNGs and also the TPFPlotter analysis to get the field of view
        information.
        @param indir: the vetting source and resources directory
        @param tic: the tic to be processed
        @param sectors_in: the sectors to be used for the given tic
        @param t0: the candidate signal first epoch
        @param period: the candidate signal period
        @param duration: the candidate signal duration
        @param depth: the candidate signal depth in ppts
        @param ffi: Whether the candidate came from FFI data
        @param run: The run where the selected light curve appeared
        @param curve: The selected light curve in the given run
        @return: the given tic
        """
        logging.info("Running TESS Point")
        sectors_all, ra, dec = LATTEutils.tess_point(indir, tic)
        try:
            sectors = list(set(sectors_in) & set(sectors_all))
            if len(sectors) == 0:
                logging.info("The target was not observed in the sector(s) you stated ({}). \
                        Therefore take all sectors that it was observed in: {}".format(sectors, sectors_all))
                sectors = sectors_all
        except:
            sectors = sectors_all
        logging.info("Downloading LATTE data")
        sectors = np.sort(sectors)
        if not ffi:
            alltime, allflux, allflux_err, all_md, alltimebinned, allfluxbinned, allx1, allx2, ally1, ally2, alltime12, allfbkg, start_sec, end_sec, in_sec, tessmag, teff, srad = LATTEutils.download_data(
                indir, sectors, tic)
        else:
            alltime_list, allflux, allflux_small, allflux_flat, all_md, allfbkg, allfbkg_t, start_sec, end_sec, in_sec, X1_list, X4_list, apmask_list, arrshape_list, tpf_filt_list, t_list, bkg_list, tpf_list = LATTEutils.download_data_FFI(indir, sectors, syspath, sectors_all, tic, True)
            srad = "-"
            tessmag = "-"
            teff = "-"
            alltime = alltime_list
        simple = False
        BLS = False
        model = False
        save = True
        DV = True
        lc_file = "/" + str(run) + "/lc_" + str(curve) + ".csv"
        lc_file = self.object_dir + lc_file
        lc = pd.read_csv(lc_file, header=0)
        time, flux, flux_err = lc["#time"].values, lc["flux"].values, lc["flux_err"].values
        lc_len = len(time)
        zeros_lc = np.zeros(lc_len)
        logging.info("Preparing folded light curves for target")
        lc = TessLightCurve(time=time, flux=flux, flux_err=flux_err, quality=zeros_lc)
        lc.extra_columns = []
        self.plot_folded_curve(self.data_dir, "TIC " + tic, lc, period, t0, duration, depth / 1000)
        transit_list = []
        last_time = alltime[len(alltime) - 1]
        num_of_transits = int(ceil(((last_time - t0) / period)))
        transit_lists = t0 + period * range(0, num_of_transits)
        time_as_array = np.array(alltime)
        transits_in_data = [time_as_array[(transit > time_as_array - 0.5) & (transit < time_as_array + 0.5)] for transit in transit_lists]
        transit_lists = transit_lists[[len(transits_in_data_set) > 0 for transits_in_data_set in transits_in_data]]
        transit_lists = [transit_lists[x:x + 3] for x in range(0, len(transit_lists), 3)]
        for index, transit_list in enumerate(transit_lists):
            transit_results_dir = self.data_dir + "/" + str(index)
            if not os.path.exists(transit_results_dir):
                os.mkdir(transit_results_dir)
            Vetter.plot_single_transits(transit_results_dir, "TIC " + str(tic), lc, transit_list, duration, depth / 1000)
            logging.info("Brewing LATTE data for transits at T0s: %s", str(transit_list))
            try:
                if not ffi:
                    LATTEbrew.brew_LATTE(tic, indir, syspath, transit_list, simple, BLS, model, save, DV, sectors,
                                         sectors_all,
                                         alltime, allflux, allflux_err, all_md, alltimebinned, allfluxbinned, allx1, allx2,
                                         ally1, ally2, alltime12, allfbkg, start_sec, end_sec, in_sec, tessmag, teff, srad, ra,
                                         dec, self.args)
                else:
                    LATTEbrew.brew_LATTE_FFI(tic, indir, syspath, transit_list, simple, BLS, model, save, DV, sectors,
                                             sectors_all, alltime, allflux_flat, allflux_small, allflux, all_md, allfbkg,
                                             allfbkg_t, start_sec, end_sec, in_sec, X1_list, X4_list, apmask_list,
                                             arrshape_list, tpf_filt_list, t_list, bkg_list, tpf_list, ra, dec, self.args)
                # LATTE_DV.LATTE_DV(tic, indir, syspath, transit_list, sectors_all, simple, BLS, model, save, DV, sectors,
                #                      sectors_all,
                #                      alltime, allflux, allflux_err, all_md, alltimebinned, allfluxbinned, allx1, allx2,
                #                      ally1, ally2, alltime12, allfbkg, start_sec, end_sec, in_sec, tessmag, teff, srad, ra,
                #                      dec, self.args)
                tp_downloaded = True
                copy_tree(vetter.latte_dir + "/" + tic, transit_results_dir)
                shutil.rmtree(vetter.latte_dir + "/" + tic, ignore_errors=True)
            except Exception as e:
                traceback.print_exc()
                # see if it made any plots - often it just fails on the TPs as they are very large
                if exists("{}/{}/{}_fullLC_md.png".format(indir, tic, tic)):
                    logging.warning("couldn't download TP but continue anyway")
                    tp_downloaded = False
                    shutil.move(vetter.latte_dir + "/" + tic, transit_results_dir)
                else:
                    continue
        # check again whether the TPs downloaded - depending on where the code failed it might still have worked.
        if exists("{}/{}/{}_aperture_size.png".format(indir, tic, tic)):
            tp_downloaded = True
        else:
            tp_downloaded = False
            logging.warn("code ran but no TP -- continue anyway")
        # -------------
        # check whether it's a TCE or a TOI

        # TCE -----
        lc_dv = np.genfromtxt('{}/data/tesscurl_sector_all_dv.sh'.format(indir), dtype=str)
        TCE_links = []
        for i in lc_dv:
            if str(tic) in str(i[6]):
                TCE_links.append(i[6])
        if len(TCE_links) == 0:
            TCE = " - "
            TCE = False
        else:
            TCE_links = np.sort(TCE_links)
            TCE_link = TCE_links[0]  # this link should allow you to acess the MAST DV report
            TCE = True
        # TOI -----TOI_planets = pd.read_csv('{}/data/TOI_list.txt'.format(indir), comment="#")
        # TOIpl = TOI_planets.loc[TOI_planets['TIC'] == float(tic)]
        #TOI = False
        # TODO check why TOI is useful
        # else:
        #     TOI = True
        #     TOI_name = (float(TOIpl["Full TOI ID"]))
        # -------------
        # return the tic so that it can be stored in the manifest to keep track of which files have already been produced
        # and to be able to skip the ones that have already been processed if the code has to be restarted.
        mnd = {}
        mnd['id'] = tic
        mnd['MarkedTransits'] = transit_list
        mnd['Sectors'] = sectors_all
        mnd['RA'] = ra
        mnd['DEC'] = dec
        mnd['SolarRad'] = srad
        mnd['TMag'] = tessmag
        mnd['Teff'] = teff
        mnd['thissector'] = sectors
        # make empty fields for the test to be checked
        mnd['TOI'] = " "
        if TCE == True:
            mnd['TCE'] = "Yes"
            mnd['TCE_link'] = TCE_link
        else:
            mnd['TCE'] = " "
            mnd['TCE_link'] = " "
        mnd['EB'] = " "
        mnd['Systematics'] = " "
        mnd['TransitShape'] = " "
        mnd['BackgroundFlux'] = " "
        mnd['CentroidsPositions'] = " "
        mnd['MomentumDumps'] = " "
        mnd['ApertureSize'] = " "
        mnd['InoutFlux'] = " "
        mnd['Keep'] = " "
        mnd['Comment'] = " "
        mnd['starttime'] = np.nanmin(alltime) if not isinstance(alltime, str) else "-"
        return mnd

    @staticmethod
    def plot_single_transits(file_dir, id, lc, transit_times, duration, depth):
        """
        Plots the phase-folded curve of the candidate for period, 2 * period and period / 2.
        @param file_dir: the directory to store the plot
        @param id: the target id
        @param lc: the input light curve with the data
        @param transit_times: the single transits T0s
        @param duration: the transit duration
        @param depth: the transit depth
        """
        duration = duration / 60 / 24
        figsize = (8, 8)  # x,y
        rows = len(transit_times)
        cols = 1
        fig, axs = plt.subplots(rows, cols, figsize=figsize, constrained_layout=True)
        for index in np.arange(0, len(transit_times)):
            current_axs = axs[index] if isinstance(axs, np.ndarray) or isinstance(axs, list) else axs
            Vetter.plot_single_transit(id, current_axs, lc, transit_times[index], depth, duration)
        plt.savefig(file_dir + "/single_transits.png", dpi=200)
        fig.clf()
        plt.close(fig)

    @staticmethod
    def plot_single_transit(id, axs, lc, transit_time, depth, duration):
        sort_args = np.argsort(lc.time.value)
        time = lc.time.value[sort_args]
        flux = lc.flux.value[sort_args]
        flux_err = lc.flux_err.value[sort_args]
        folded_plot_range = duration * 2
        folded_phase_zoom_mask = np.where((time > transit_time - folded_plot_range) &
                                          (time < transit_time + folded_plot_range))
        folded_phase = time[folded_phase_zoom_mask]
        folded_y = flux[folded_phase_zoom_mask]
        if len(folded_phase) > 0:
            in_transit = (folded_phase > transit_time - duration / 2) & (folded_phase < transit_time + duration / 2)
            in_transit_center = (np.abs(folded_phase - transit_time)).argmin()
            model_time, model_flux = Vetter.get_transit_model(duration, transit_time,
                                                  (transit_time - folded_plot_range, transit_time + folded_plot_range),
                                                  depth)
            axs.plot(model_time, model_flux, color="red")
            axs.scatter(folded_phase[~in_transit], folded_y[~in_transit], color="gray")
            axs.scatter(folded_phase[in_transit], folded_y[in_transit], color="darkorange")
            axs.set_xlim([transit_time - folded_plot_range, transit_time + folded_plot_range])
            axs.set_title(str(id) + " Single Transit at T={:.2f}d".format(transit_time))
            #axs.set_ylim([1 - 3 * depth, 1 + 3 * depth])
            axs.set_xlim([transit_time - folded_plot_range, transit_time + folded_plot_range])
            axs.set_xlabel("Time (d)")
            axs.set_ylabel("Flux norm.")
            logging.info("Processed single transit plot for T0=%.2f", transit_time)
        else:
            logging.info("Not plotting single transit for T0=%.2f as the data is empty", transit_time)

    @staticmethod
    def plot_folded_curve(file_dir, id, lc, period, epoch, duration, depth):
        """
        Plots the phase-folded curve of the candidate for period, 2 * period and period / 2.
        @param file_dir: the directory to store the plot
        @param id: the target id
        @param period: the transit period
        @param epoch: the transit epoch
        @param duration: the transit duration
        @param depth: the transit depth
        """
        duration = duration / 60 / 24
        figsize = (16, 16)  # x,y
        rows = 3
        cols = 2
        fig, axs = plt.subplots(rows, cols, figsize=figsize, constrained_layout=True)
        logging.info("Preparing folded light curves for target")
        #TODO bins = None for FFI
        bins = 100
        Vetter.compute_phased_values_and_fill_plot(id, axs[0][0], lc, period, epoch + period / 2, depth, duration,
                                                   bins=bins)
        Vetter.compute_phased_values_and_fill_plot(id, axs[0][1], lc, period, epoch, depth, duration, bins=bins)
        period = 2 * period
        Vetter.compute_phased_values_and_fill_plot(id, axs[1][0], lc, period, epoch + period / 2, depth, duration,
                                                   bins=bins)
        Vetter.compute_phased_values_and_fill_plot(id, axs[1][1], lc, period, epoch, depth, duration, bins=bins)
        period = period / 4
        Vetter.compute_phased_values_and_fill_plot(id, axs[2][0], lc, period, epoch + period / 2, depth, duration,
                                                   bins=bins)
        Vetter.compute_phased_values_and_fill_plot(id, axs[2][1], lc, period, epoch, depth, duration, bins=bins)
        plt.savefig(file_dir + "/odd_even_folded_curves.png", dpi=200)
        fig.clf()
        plt.close(fig)

    @staticmethod
    def compute_phased_values_and_fill_plot(id, axs, lc, period, epoch, depth, duration, range=5, bins=None):
        """
        Phase-folds the input light curve and plots it centered in the given epoch
        @param id: the candidate name
        @param axs: the plot axis to be drawn
        @param lc: the lightkurve object containing the data
        @param period: the period for the phase-folding
        @param epoch: the epoch to center the fold
        @param depth: the transit depth
        @param duration: the transit duration
        @param range: the range to be used from the midtransit time in half-duration units.
        @param bins: the number of bins
        @return: the drawn axis and the computed bins
        """
        time = foldedleastsquares.core.fold(lc.time.value, period, epoch)
        axs.scatter(time, lc.flux.value, 2, color="blue", alpha=0.1)
        sort_args = np.argsort(time)
        time = time[sort_args]
        flux = lc.flux.value[sort_args]
        flux_err = lc.flux_err.value[sort_args]
        half_duration_phase = duration / 2 / period
        folded_plot_range = half_duration_phase * range
        folded_plot_range = folded_plot_range if folded_plot_range < 0.5 else 0.5
        folded_phase_zoom_mask = np.where((time > 0.5 - folded_plot_range) &
                                          (time < 0.5 + folded_plot_range))
        folded_phase = time[folded_phase_zoom_mask]
        folded_y = flux[folded_phase_zoom_mask]
        folded_y_err = flux_err[folded_phase_zoom_mask]
        axs.set_xlim([0.5 - folded_plot_range, 0.5 + folded_plot_range])
        # TODO if FFI no binning
        if bins is not None:
            bin_means, bin_edges, binnumber = stats.binned_statistic(folded_phase, folded_y,
                                                                     statistic='mean', bins=bins)
            bin_width = (bin_edges[1] - bin_edges[0])
            bin_centers = bin_edges[1:] - bin_width / 2
            bin_stds, _, _ = stats.binned_statistic(folded_phase, folded_y, statistic='std', bins=bins)
            axs.errorbar(bin_centers, bin_means, yerr=bin_stds / 2, xerr=bin_width / 2, marker='o', markersize=4,
                         color='darkorange', alpha=1, linestyle='none')
        else:
            bin_centers = folded_phase
            bin_means = folded_y
            bin_stds = folded_y_err * 2
        in_transit = (folded_phase > 0.5 - half_duration_phase) & (folded_phase < 0.5 + half_duration_phase)
        in_transit_center = (np.abs(folded_phase - 0.5)).argmin()
        model_time, model_flux = Vetter.get_transit_model(half_duration_phase * 2, 0.5, (0, 1.0), depth)
        axs.plot(model_time, model_flux, color="red")
        axs.set_title(str(id) + " Folded Curve with P={:.2f}d and T0={:.2f}".format(period, epoch))
        axs.set_xlabel("Time (d)")
        axs.set_ylabel("Flux norm.")
        #axs.set_ylim([1 - 3 * depth, 1 + 3 * depth])
        logging.info("Processed phase-folded plot for P=%.2f and T0=%.2f", period, epoch)
        return axs, bin_centers, bin_means, bin_stds / 2

    @staticmethod
    #TODO build model from selected transit_template
    def get_transit_model(duration, t0, start_end, depth):
        model_len = 10000
        t = np.linspace(-6, 6, model_len)
        ma = batman.TransitParams()
        ma.t0 = 0  # time of inferior conjunction
        ma.per = 365  # orbital period, use Earth as a reference
        ma.rp = 0.1  # planet radius (in units of stellar radii)
        ma.a = 50  # semi-major axis (in units of stellar radii)
        ma.inc = 90  # orbital inclination (in degrees)
        ma.ecc = 0  # eccentricity
        ma.w = 0  # longitude of periastron (in degrees)
        ma.u = [0.4804, 0.1867]  # limb darkening coefficients
        ma.limb_dark = "quadratic"  # limb darkening model
        m = batman.TransitModel(ma, t)  # initializes model
        model = m.light_curve(ma)  # calculates light curve
        model_time = np.linspace(start_end[0], start_end[1], model_len)
        model_intransit = np.argwhere(model < 1)[:, 0]
        in_transit_indexes = np.where((model_time > t0 - duration / 2) & (model_time < t0 + duration / 2))[0]
        model_time_in_transit = model_time[in_transit_indexes]
        scaled_intransit = np.interp(
            np.linspace(model_time_in_transit[0], model_time_in_transit[-1], len(in_transit_indexes)),
            model_time[model_intransit], model[model_intransit])
        model = np.full((model_len), 1.0)
        model[in_transit_indexes] = scaled_intransit
        model[model < 1] = 1 - ((1 - model[model < 1]) * depth / (1 - np.min(model)))
        return model_time, model

    def vetting(self, candidate, cpus):
        """
        Performs the LATTE vetting procedures
        @param candidate: the candidate dataframe containing id, period, t0, transits and sectors data.
        @param cpus: the number of cpus to be used. This parameter is of no use yet.
        """
        indir, df, TIC_wanted, ffi = self.__prepare(candidate)
        for tic in TIC_wanted:
            # check the existing manifest to see if I've processed this file!
            manifest_table = pd.read_csv('{}/manifest.csv'.format(str(indir)))
            # get the transit time list
            period = df.loc[df['id'] == tic]['period'].iloc[0]
            t0 = df.loc[df['id'] == tic]['t0'].iloc[0]
            duration = df.loc[df['id'] == tic]['duration'].iloc[0]
            depth = df.loc[df['id'] == tic]['depth'].iloc[0]
            ffi = df.loc[df['id'] == tic]['ffi'].iloc[0]
            run = int(candidate.loc[candidate['id'] == tic]['number'].iloc[0])
            curve = int(candidate.loc[candidate['id'] == tic]['curve'].iloc[0])
            logging.info("------------------")
            logging.info("Candidate info")
            logging.info("------------------")
            logging.info("Period (d): %.2f", period)
            logging.info("Epoch (d): %.2f", t0)
            logging.info("Duration (min): %.2f", duration)
            logging.info("Depth (ppt): %.2f", depth)
            logging.info("FFI: %s", ffi)
            logging.info("Run: %.0f", run)
            logging.info("Detrend curve: %.0f", curve)
            try:
                sectors_in = ast.literal_eval(str(((df.loc[df['id'] == tic]['sectors']).values)[0]))
                if (type(sectors_in) == int) or (type(sectors_in) == float):
                    sectors = [sectors_in]
                else:
                    sectors = list(sectors_in)
            except:
                sectors = [0]
            index = 0
            vetting_dir = self.data_dir + "/vetting_" + str(index)
            while os.path.exists(vetting_dir) or os.path.isdir(vetting_dir):
                vetting_dir = self.data_dir + "/vetting_" + str(index)
                index = index + 1
            os.mkdir(vetting_dir)
            self.data_dir = vetting_dir
            ra = None
            dec = None
            try:
                res = self.__process(indir, tic, sectors, t0, period, duration, depth, ffi, run, curve)
                ra = res['RA']
                dec = res['DEC']
                if res['id'] == -99:
                    logging.error('something went wrong with the LATTE results')
            except Exception as e:
                traceback.print_exc()
                try:
                    sectors_all, ra, dec = LATTEutils.tess_point(indir, tic)
                except Exception as e1:
                    traceback.print_exc()

    @staticmethod
    def vetting_field_of_view(indir, mission, tic, cadence, ra, dec, sectors, source, apertures):
        """
        Runs TPFPlotter to get field of view data.
        @param indir: the data source directory
        @param mission: the mission of the target
        @param tic: the target id
        @param cadence: the exposure time between measurements in seconds
        @param ra: the right ascension of the target
        @param dec: the declination of the target
        @param sectors: the sectors where the target was observed
        @param source: the source where the aperture was generated [tpf, tesscut, eleanor]
        @param apertures: a dict mapping sectors to boolean apertures
        @return: the directory where resulting data is stored
        """
        try:
            maglim = 6
            sectors_search = None if sectors is not None and len(sectors) == 0 else sectors
            logging.info("Preparing target pixel files for field of view plots")
            ra_str = str(ra)
            dec_str = "+" + str(dec) if dec >= 0 else str(dec)
            coords_str = ra_str + " " + dec_str
            if mission != "TESS":
                return
            target_title = "TIC " + str(tic)
            #TODO use retrieval method depending on source parameter
            if cadence > 120:
                tpf_source = lightkurve.search_tesscut(target_title, sector=sectors_search)
                if tpf_source is None or len(tpf_source) == 0:
                    tpf_source = lightkurve.search_tesscut(coords_str, sector=sectors_search)
                    target_title = "RA={:.4f},DEC={:.4f}".format(ra, dec)
            else:
                tpf_source = lightkurve.search_targetpixelfile(target_title, sector=sectors_search, author="SPOC",
                                                               cadence=cadence)
            save_dir = indir
            for i in range(0, len(tpf_source)):
                tpf = tpf_source[i].download(cutout_size=(CUTOUT_SIZE, CUTOUT_SIZE))
                row = tpf.row
                column = tpf.column
                plt.close()
                fig = plt.figure(figsize=(6.93, 5.5))
                gs = gridspec.GridSpec(1, 3, height_ratios=[1], width_ratios=[1, 0.05, 0.01])
                gs.update(left=0.05, right=0.95, bottom=0.12, top=0.95, wspace=0.01, hspace=0.03)
                ax1 = plt.subplot(gs[0, 0])
                # TPF plot
                mean_tpf = np.mean(tpf.flux.value, axis=0)
                nx, ny = np.shape(mean_tpf)
                norm = ImageNormalize(stretch=stretching.LogStretch())
                division = np.int(np.log10(np.nanmax(tpf.flux.value)))
                splot = plt.imshow(np.nanmean(tpf.flux, axis=0) / 10 ** division, norm=norm, cmap="viridis",\
                                   extent=[column, column + ny, row, row + nx], origin='lower', zorder=0)
                aperture = apertures[tpf.sector]
                aperture = aperture if isinstance(aperture, np.ndarray) else np.array(aperture)
                maskcolor = 'salmon'
                logging.info("    --> Using SHERLOCK aperture for sector %s...", tpf.sector)
                if aperture is not None:
                    for pixels in aperture:
                        ax1.add_patch(patches.Rectangle((pixels[0], pixels[1]),
                                                        1, 1, color=maskcolor, fill=True, alpha=0.4))
                        ax1.add_patch(patches.Rectangle((pixels[0], pixels[1]),
                                                        1, 1, color=maskcolor, fill=False, alpha=1, lw=2))
                # Gaia sources
                gaia_id, mag = tpfplotter.get_gaia_data(ra, dec)
                r, res = tpfplotter.add_gaia_figure_elements(tpf, magnitude_limit=mag + np.float(maglim), targ_mag=mag)
                x, y, gaiamags = r
                x, y, gaiamags = np.array(x) + 0.5, np.array(y) + 0.5, np.array(gaiamags)
                size = 128.0 / 2 ** ((gaiamags - mag))
                plt.scatter(x, y, s=size, c='red', alpha=0.6, edgecolor=None, zorder=10)
                # Gaia source for the target
                this = np.where(np.array(res['Source']) == int(gaia_id))[0]
                plt.scatter(x[this], y[this], marker='x', c='white', s=32, zorder=11)
                # Legend
                add = 0
                if np.int(maglim) % 2 != 0:
                    add = 1
                maxmag = np.int(maglim) + add
                legend_mags = np.linspace(-2, maxmag, np.int((maxmag + 2) / 2 + 1))
                fake_sizes = mag + legend_mags  # np.array([mag-2,mag,mag+2,mag+5, mag+8])
                for f in fake_sizes:
                    size = 128.0 / 2 ** ((f - mag))
                    plt.scatter(0, 0, s=size, c='red', alpha=0.6, edgecolor=None, zorder=10,
                                label=r'$\Delta m=$ ' + str(np.int(f - mag)))
                ax1.legend(fancybox=True, framealpha=0.7)
                # Source labels
                dist = np.sqrt((x - x[this]) ** 2 + (y - y[this]) ** 2)
                dsort = np.argsort(dist)
                for d, elem in enumerate(dsort):
                    if dist[elem] < 6:
                        plt.text(x[elem] + 0.1, y[elem] + 0.1, str(d + 1), color='white', zorder=100)
                # Orientation arrows
                tpfplotter.plot_orientation(tpf)
                # Labels and titles
                plt.xlim(column, column + ny)
                plt.ylim(row, row + nx)
                plt.xlabel('Pixel Column Number', fontsize=16)
                plt.ylabel('Pixel Row Number', fontsize=16)
                plt.title('Coordinates ' + target_title + ' - Sector ' + str(tpf.sector),
                          fontsize=16)  # + ' - Camera '+str(tpf.camera))  #
                # Colorbar
                cbax = plt.subplot(gs[0, 1])  # Place it where it should be.
                pos1 = cbax.get_position()  # get the original position
                pos2 = [pos1.x0 - 0.05, pos1.y0, pos1.width, pos1.height]
                cbax.set_position(pos2)  # set a new position
                cb = Colorbar(ax=cbax, cmap="viridis", mappable=splot, orientation='vertical', ticklocation='right')
                plt.xticks(fontsize=14)
                exponent = r'$\times 10^' + str(division) + '$'
                cb.set_label(r'Flux ' + exponent + r' (e$^-$)', labelpad=10, fontsize=16)
                save_dir = indir
                if not os.path.exists(save_dir):
                    os.makedirs(save_dir)
                plt.savefig(save_dir + '/TPF_Gaia_' + target_title + '_S' + str(tpf.sector) + '.pdf')
                # Save Gaia sources info
                dist = np.sqrt((x - x[this]) ** 2 + (y - y[this]) ** 2)
                GaiaID = np.array(res['Source'])
                srt = np.argsort(dist)
                x, y, gaiamags, dist, GaiaID = x[srt], y[srt], gaiamags[srt], dist[srt], GaiaID[srt]
                IDs = np.arange(len(x)) + 1
                inside = np.zeros(len(x))
                for pixels in aperture:
                    xtpf, ytpf = pixels[0], pixels[1]
                    _inside = np.where((x > xtpf) & (x < xtpf + 1) &
                                       (y > ytpf) & (y < ytpf + 1))[0]
                    inside[_inside] = 1
                data = Table([IDs, GaiaID, x, y, dist, dist * 21., gaiamags, inside.astype('int')],
                             names=['# ID', 'GaiaID', 'x', 'y', 'Dist_pix', 'Dist_arcsec', 'Gmag', 'InAper'])
                ascii.write(data, save_dir + '/Gaia_' + target_title + '_S' + str(tpf.sector) + '.dat', overwrite=True)
            return save_dir
        except SystemExit:
            logging.exception("Field Of View generation tried to exit.")
        except Exception as e:
            logging.exception("Exception found when generating Field Of View plots")

if __name__ == '__main__':
    ap = ArgumentParser(description='Vetting of Sherlock objects of interest')
    ap.add_argument('--object_dir', help="If the object directory is not your current one you need to provide the "
                                         "ABSOLUTE path", required=False)
    ap.add_argument('--candidate', type=int, default=None, help="The candidate signal to be used.", required=False)
    ap.add_argument('--properties', help="The YAML file to be used as input.", required=False)
    ap.add_argument('--cpus', type=int, default=None, help="The number of CPU cores to be used.", required=False)
    ap.add_argument('--no_validate', dest='validate', action='store_false',
                    help="Whether to avoid running statistical validation")
    args = ap.parse_args()
    vetter = Vetter(args.object_dir, args.validate)
    file_dir = vetter.object_dir + "/vetting.log"
    if os.path.exists(file_dir):
        os.remove(file_dir)
    formatter = logging.Formatter('%(message)s')
    logger = logging.getLogger()
    while len(logger.handlers) > 0:
        logger.handlers.pop()
    logger.setLevel(logging.INFO)
    handler = logging.StreamHandler(sys.stdout)
    handler.setLevel(logging.INFO)
    handler.setFormatter(formatter)
    logger.addHandler(handler)
    handler = logging.FileHandler(file_dir)
    handler.setLevel(logging.INFO)
    handler.setFormatter(formatter)
    logger.addHandler(handler)
    logging.info("Starting vetting")
    star_df = pd.read_csv(vetter.object_dir + "/params_star.csv")
    if args.candidate is None:
        user_properties = yaml.load(open(args.properties), yaml.SafeLoader)
        candidate = pd.DataFrame(columns=['id', 'period', 'depth', 't0', 'sectors', 'ffi', 'number', 'lc'])
        candidate = candidate.append(user_properties, ignore_index=True)
        candidate['id'] = star_df.iloc[0]["obj_id"]
    else:
        candidate_selection = int(args.candidate)
        candidates = pd.read_csv(vetter.object_dir + "/candidates.csv")
        if candidate_selection < 1 or candidate_selection > len(candidates.index):
            raise SystemExit("User selected a wrong candidate number.")
        candidates = candidates.rename(columns={'Object Id': 'id'})
        candidate = candidates.iloc[[candidate_selection - 1]]
        candidate['number'] = [candidate_selection]
        vetter.data_dir = vetter.object_dir
        logging.info("Selected signal number " + str(candidate_selection))
    if args.cpus is None:
        cpus = multiprocessing.cpu_count() - 1
    else:
        cpus = args.cpus
    vetter.vetting(candidate, cpus)

import numpy,from scipy,import astropy,from astropy


import numpy as np
import imutils
import glob
import cv2

image = cv2.imread('COD_template.jpg')
image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
image = cv2.Canny(image, 50, 200)

cv2.imshow("image", image)
cv2.waitKey(0)
cv2.destroyAllWindows()


from copy import deepcopy
from pathlib import Path

import numpy as np
import pygeos as pg
import geopandas as gp

from analysis.lib.pygeos_util import (
    to_crs,
    to_dict,
    sjoin,
    intersection,
)
from analysis.constants import (
    INPUTS,
    INPUT_AREA_VALUES,
    URBAN_YEARS,
    DATA_CRS,
    OWNERSHIP,
    PROTECTION,
    M2_ACRES,
)
from analysis.lib.stats import (
    extract_blueprint_by_geometry,
    extract_urban_by_geometry,
    extract_slr_by_geometry,
    summarize_caribbean_by_aoi,
    summarize_chat_by_aoi,
    summarize_florida_by_aoi,
    summarize_florida_marine_by_aoi,
    summarize_gulf_hypoxia_by_aoi,
    summarize_midse_by_aoi,
    summarize_naturescape_by_aoi,
    summarize_natures_network_by_aoi,
    summarize_southatlantic_by_aoi,
)

data_dir = Path("data/inputs")
boundary_filename = data_dir / "boundaries/se_boundary.feather"
county_filename = data_dir / "boundaries/counties.feather"
ownership_filename = data_dir / "boundaries/ownership.feather"
slr_bounds_filename = data_dir / "threats/slr/slr_bounds.feather"

raster_result_funcs = {
    "fl": summarize_florida_by_aoi,
    "flm": summarize_florida_marine_by_aoi,
    "gh": summarize_gulf_hypoxia_by_aoi,
    "ms": summarize_midse_by_aoi,
    "app": summarize_naturescape_by_aoi,
    "nn": summarize_natures_network_by_aoi,
    "sa": summarize_southatlantic_by_aoi,
}

class CustomArea(object):
    def __init__(self, geometry, crs, name):
        """Initialize a custom area from a pygeos geometry.

        Parameters
        ----------
        geometry : pygeos Geometry
        crs : pyproj CRS object
        name : string
            name of custom area
        """

        self.geometry = to_crs(geometry, crs, DATA_CRS)
        self.gdf = gp.GeoDataFrame({"geometry": self.geometry}, crs=DATA_CRS)
        self.bounds = pg.total_bounds(self.geometry)
        # wrap geometry as a dict for rasterio
        self.shapes = np.asarray([to_dict(self.geometry[0])])
        self.name = name

    def get_blueprint(self):
        blueprint = extract_blueprint_by_geometry(self.shapes, bounds=self.bounds)

        if blueprint is None or blueprint["shape_mask"] == 0:
            return None

        blueprint_total = blueprint["blueprint"].sum()

        remainder = abs(blueprint["shape_mask"] - blueprint_total)
        # there are small rounding errors
        remainder = remainder if remainder >= 1 else 0

        # only pull in Blueprint inputs that are present, and flatten
        # overlapping inputs
        inputs = dict()
        has_overlapping_inputs = False
        for i, acres in enumerate(blueprint["inputs"]):
            input_ids = INPUT_AREA_VALUES[i]["id"].split(",")
            if acres > 0:
                if len(input_ids) > 1:
                    has_overlapping_inputs = True

                for input_id in input_ids:
                    if input_id not in inputs:
                        input = deepcopy(INPUTS[input_id])
                        input["acres"] = acres
                        inputs[input_id] = input
                    else:
                        inputs[input_id]["acres"] += acres

        inputs = sorted(inputs.values(), key=lambda x: x["acres"], reverse=True)

        # extract input priorities from each input present
        df = gp.GeoDataFrame(geometry=self.geometry, crs=DATA_CRS)
        aoi_acres = pg.area(self.geometry).sum() * M2_ACRES

        for entry in inputs:
            input_id = entry["id"]

            if input_id in ["okchat", "txchat"]:
                state = input_id[:2]
                chat_results = summarize_chat_by_aoi(df, state, aoi_acres)

                if chat_results is not None:
                    entry.update(chat_results)

                continue

            if input_id == "car":
                caribbean_results = summarize_caribbean_by_aoi(df, aoi_acres, aoi_acres)
                entry.update(caribbean_results)

                continue

            # Remaining inputs are raster-based
            results_func = raster_result_funcs.get(input_id, None)
            if not results_func:
                print(f"WARNING: missing AOI summary func for {input_id}")
                continue

            raster_results = results_func(
                self.shapes, self.bounds, outside_se_acres=remainder
            )

            if raster_results is not None:
                entry.update(raster_results)
            else:
                print("Raster results are none for", input_id)

        results = {
            "analysis_acres": blueprint["shape_mask"],
            "analysis_remainder": remainder,
            "blueprint": blueprint["blueprint"].tolist(),
            "blueprint_total": blueprint_total,
            "inputs": inputs,
            "input_ids": [i["id"] for i in inputs],
            "has_overlapping_inputs": has_overlapping_inputs,
        }

        return results

    def get_urban(self):
        """Extract current and projected urbanization.

        Returns
        -------
        dict
            {
                "urban_acres": <urban analysis acres>,
                "urban": <current urban>,
                "proj_urban": [<urbanization 2020>, ..., <urbanization 2060>]
            }
        """
        urban = extract_urban_by_geometry(self.shapes, bounds=self.bounds)

        if urban is None or urban["shape_mask"] == 0:
            return None

        # only keep through 2060 for reports, and only if there is projected
        # urbanization
        proj_urban = [urban[year] for year in URBAN_YEARS[:5]]
        if not sum(proj_urban):
            return None

        return {
            "urban_acres": urban["shape_mask"],
            "urban": urban["urban"],
            "proj_urban": proj_urban,
        }

    def get_slr(self):
        """Extract SLR for any geometries that overlap bounds where SLR is available

        Returns
        -------
        dict
            {"slr_acres": <acres>, "slr": [<slr_0ft>, <slr_1ft>, ..., <slr_6ft>]}
        """
        slr_bounds = gp.read_feather(slr_bounds_filename).geometry.values.data[0]
        ix = pg.intersects(self.geometry, slr_bounds)

        if not ix.sum():
            # No overlap
            return None

        # only extract SLR where there are overlaps
        slr_results = extract_slr_by_geometry(
            self.shapes[ix], bounds=pg.total_bounds(self.geometry[ix])
        )
        # None only if no shape mask
        if slr_results is None:
            return None

        slr = [slr_results[i] for i in range(7)]

        return {"slr_acres": slr_results["shape_mask"], "slr": slr}

    def get_counties(self):
        """Get county and state names that overlap this area.

        Returns
        -------
        dict
            {"counties": [
                {"FIPS": <FIPS>, "state": <state name>, "county": <county_name>},
                ...
            ]
        """
        counties = gp.read_feather(county_filename)[
            ["geometry", "FIPS", "state", "county"]
        ]

        df = (
            sjoin(self.gdf, counties)[["FIPS", "state", "county"]]
            .reset_index(drop=True)
            .sort_values(by=["state", "county"])
        )

        if not len(df):
            return None

        return {"counties": df.to_dict(orient="records")}

    def get_ownership(self):
        """Get ownership and protection levels and other statistics for this area

        Returns
        -------
        dict
            {
                "ownership": [
                    {
                        "label": <ownership type label>,
                        "acres": <acres of overlap>
                    }
                ],
                "protection": [
                    {
                        "label": <protection type label>,
                        "acres": <acres of overlap>
                    }
                ],
                "protected_areas" [<top 25 protected area names and areas>],
                "num_protected_areas": <total number protected areas>
            }
        """
        ownership = gp.read_feather(ownership_filename)
        df = intersection(self.gdf, ownership)

        if not len(df):
            return None

        df["acres"] = pg.area(df.geometry_right.values.data) * M2_ACRES
        df = df.loc[df.acres > 0].copy()

        if not len(df):
            return None

        results = dict()

        by_owner = (
            df[["Own_Type", "acres"]]
            .groupby(by="Own_Type")
            .acres.sum()
            .astype("float32")
            .to_dict()
        )
        # use the native order of OWNERSHIP to drive order of results
        results["ownership"] = [
            {"label": value["label"], "acres": by_owner[key]}
            for key, value in OWNERSHIP.items()
            if key in by_owner
        ]

        by_protection = (
            df[["GAP_Sts", "acres"]]
            .groupby(by="GAP_Sts")
            .acres.sum()
            .astype("float32")
            .to_dict()
        )
        # use the native order of PROTECTION to drive order of results
        results["protection"] = [
            {"label": value["label"], "acres": by_protection[key]}
            for key, value in PROTECTION.items()
            if key in by_protection
        ]

        by_area = (
            df[["Loc_Nm", "Loc_Own", "acres"]]
            .groupby(by=[df.index.get_level_values(0), "Loc_Nm", "Loc_Own"])
            .acres.sum()
            .astype("float32")
            .round()
            .reset_index()
            .rename(columns={"level_0": "id", "Loc_Nm": "name", "Loc_Own": "owner"})
            .sort_values(by="acres", ascending=False)
        )
        # drop very small areas, these are not helpful
        by_area = by_area.loc[by_area.acres >= 1].copy()

        results["protected_areas"] = by_area.head(25).to_dict(orient="records")
        results["num_protected_areas"] = len(by_area)

        return results

    def get_results(self):
        se_bnd = gp.read_feather(boundary_filename)

        # if area of interest does not intersect SE region boundary,
        # there will be no results
        if not pg.intersects(self.geometry, se_bnd.geometry.values.data).max():
            return None

        results = {
            "type": "",
            "acres": pg.area(self.geometry).sum() * M2_ACRES,
            "name": self.name,
        }

        blueprint_results = self.get_blueprint()
        if blueprint_results is None:
            return None

        results.update(blueprint_results)

        urban_results = self.get_urban()
        if urban_results is not None:
            results.update(urban_results)

        slr_results = self.get_slr()
        if slr_results is not None:
            results.update(slr_results)

        ownership_results = self.get_ownership()
        if ownership_results is not None:
            results.update(ownership_results)

        county_results = self.get_counties()
        if county_results is not None:
            results.update(county_results)

        return results


'''
init: Sept-2019
Author: OO
Topic: Proximal Policy Optimization Implementation for Binary action-space env.
Issues: PPO is mostly only useful when there is a distributed worker architecture
with a central node issuing π_old evaluations to compare with π_t for the same {(a_t|s_t)}_t.
Goals: All restricted to Policy Gradient Methods...
    - `EmbodiedAgent` implements: (as base class)
        - tripartite algo split
        - embodiment concept
        - play() construct for generating experience tuples
    - `EmbodiedAgent_<Algo>` implements:
        - <Algo>'s training procedures
        - overrides work() fxn in base class
To Do:
    - Implement online training version
'''
import numpy as np
import tensorflow as tf
# from minoritygame.minority_env import MinorityGame1vN_env
from embodied import *
from embodied_misc import _zdim_, _sched_win_, cyc_learning_spd, policy_entropy
from embodied_misc import calc_V_TD_target, policy_entropy, bernoulli_H, bernoulli_LL

# from embodied_arch.embodied import *

tboard_path = "./log"
agent_name = "embodied_PPO"
__version__ = "0.0.1"
_DEBUG_ = False

# model: var defaults
_default_reports_ = ['Perf/Recent Reward',
                     'Losses/Policy LL', 'Losses/Entropy',
                     'Norms/Grad Norm', 'Norms/Var Norm']

_every_ = 100
_eps_ = 1.e-4
_clip_default_ = 0.2

_g_thresh_ = 20.0  # gradient clip threshold
_ent_decay_ = 5e-3
_zdim_ = 16
(lrp, lrv) = (1e-3, 1e-2)  # learning rates

# envs: var defaults
(na_, m_, s_, p_) = (33, 3, 4, 0.5)
_max_len_ = 200
_s_size_, _a_size_ = (4, 1)

class EmbodiedAgentPPO(EmbodiedAgent):
    def __init__(self, name=agent_name,
                 env_=gym.make('CartPole-v0'),
                 alpha_p=5e-2, alpha_v=1e-1, clipping=_clip_default_,
                 latentDim=_zdim_,
                 space_size=(_s_size_, _a_size_),
                 sensorium=SensoriumNetworkTemplate,
                 actorNN=ActionPolicyNetwork,
                 valueNN=ValueNetwork,
                 recover=None, _every_=_every_,
                 max_episode_length=_max_len_
                 ):
        super().__init__(name=name, env_=env_,
                         latentDim=latentDim, space_size=space_size,
                         sensorium=sensorium, actorNN=actorNN, valueNN=valueNN,
                         _every_=_every_, max_episode_length=max_episode_length
                         )

        self.optimizer_p = tf.train.AdamOptimizer(learning_rate=lrp)
        self.optimizer_v = tf.train.AdamOptimizer(learning_rate=lrv)
        self.clip_range = clipping

        self.report_labels = ['Perf/Recent Reward',
                              'Losses/Policy LL', 'Losses/Value Fxn', 'Losses/Entropy']
        self.last_good_model = recover

        with tf.variable_scope(self.name):
            # AC/PPO-specific placeholders: ln π, delta-Advantages, ratio
            self.old_lnPi_t = tf.placeholder(dtype=tf.float32, shape=[None, 1],
                                             name="old_ln_pi")

            # Intermediate variables
            # Evaluate ln π(a_t|s_t) for bernoulli action space output
            # Taking advantage of built-in cross_entropy fxn (w/ numeric guards). NB: need -1*...
            self.lnPi_t = - tf.nn.sigmoid_cross_entropy_with_logits(
                logits=self.a_logit,  ## a_t|s_t
                labels=self.actions_At,  ## use tf.one_hot for m-ary action spaces
                name="ln_pi"
            )
            self.policy_LL = tf.reduce_mean(self.lnPi_t)  # Report vars
            self.ratio = tf.exp(
                self.lnPi_t - self.old_lnPi_t,
                name="pi_ratio")  # holds: ratio bw new/old π
            self.Advs_t = self.returns_Gt - self.value

            self.clipped_ratio = tf.clip_by_value(
                self.ratio,
                1.0 - self.clip_range, 1.0 + self.clip_range,
                name="clipped_ratio"
            )
            self.policy_obj = tf.minimum(
                self.ratio * tf.stop_gradient(self.Advs_t),
                self.clipped_ratio * tf.stop_gradient(self.Advs_t),
                name="policy_objective"
            )  # replaces standard ∂_t * ln π_t PG-objective
            self.entropy = tf.clip_by_value(
                tf.reduce_mean(bernoulli_H(self.a_prob)),
                _eps_, 100.
            )
            # Losses
            self.ploss = tf.reduce_mean(self.policy_obj)
            self.vloss = 0.5 * tf.reduce_mean(
                tf.square(self.returns_Gt - self.value)
            )

            # Separate out: Training Vars
            self.policy_vars = [v for v in tf.trainable_variables()
                                if "action" in v.name or "sensorium" in v.name]
            self.value_vars = [v for v in tf.trainable_variables()
                               if "value_fxn" in v.name or "sensorium" in v.name]

            # Training Ops for p+v nets
            self.p_grads = self.optimizer_p.compute_gradients(
                loss=(-alpha_p) * (self.ploss + _ent_decay_ * self.entropy),
                var_list=self.policy_vars
            )
            self.p_train = self.optimizer_p.apply_gradients(self.p_grads)
            # self.optimizer_p.minimize(-alpha_p * self.ploss)
            self.v_train = self.optimizer_v.minimize(
                alpha_v * self.vloss,
                var_list=self.value_vars
            )
        return

    def act(self, state, sess):
        """Returns policy net sample action (in {0,1})"""
        probs = sess.run(self.a_prob, {self.states_St: np.expand_dims(state.flatten(), axis=0)})
        a_t = 1 * (np.random.rand() < probs)
        return a_t.squeeze()

    def train(self, sess, gamma=0.99, bootstrap_value=0.0, window=None):
        # parse buffer
        states = np.stack(self.episode_buffer['states'])
        actions = np.stack(self.episode_buffer['actions'])
        rewards = np.stack(self.episode_buffer['rewards'])
        next_states = np.stack(self.episode_buffer['next_states'])

        # generate state-value estimates; goes into delta_t
        rwds = rewards.ravel()
        vals = sess.run(self.value,
                        feed_dict={self.states_St: np.vstack(states)}).ravel()
        old_pis = sess.run(self.lnPi_t,
                           feed_dict={
                               self.actions_At: np.vstack(actions),
                               self.states_St: np.vstack(states)
                           }).ravel()  # cheating w/ old_pi = pi for now.
        Gt_TD = np.squeeze(calc_V_TD_target(rwds, vals,
                                            gamma=gamma, bootstrap_value=bootstrap_value))
        feed_dict = {
            self.states_St: np.vstack(states),
            self.actions_At: np.vstack(actions),
            self.returns_Gt: np.vstack(Gt_TD),
            self.old_lnPi_t: np.vstack(old_pis)
        }
        # Train
        sess.run([self.v_train, self.p_train], feed_dict=feed_dict)
        # Generate performance statistics to periodically save
        p_ll, v_l, p_ent = sess.run(
            [self.ploss, self.vloss, self.entropy],
            feed_dict=feed_dict
        )
        return p_ll, v_l, p_ent

    def pretrainV(self, sess, gamma=0.95):
        # parse buffer
        states = np.stack(self.episode_buffer['states'])
        rewards = np.stack(self.episode_buffer['rewards'])

        # generate state-value estimates; goes into delta_t
        rwds = rewards.ravel()
        vals = sess.run(self.value, feed_dict={self.states_St: np.vstack(states)}).ravel()
        Gt_TD = np.squeeze(calc_V_TD_target(rwds, vals, gamma=gamma))
        feed_dict = {
            self.states_St: np.vstack(states),
            self.returns_Gt: np.vstack(Gt_TD)
        }
        # Train & Record v_loss
        sess.run([self.v_train], feed_dict=feed_dict)
        v_l = sess.run(self.vloss, feed_dict=feed_dict)
        return v_l