# from __future__ import print_function, absolute_import, division
import logging
import multiprocessing
import traceback

import PIL
import batman
import foldedleastsquares
import lcbuilder
import lightkurve
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec
import wotan
import yaml
from lcbuilder.constants import CUTOUT_SIZE
from lcbuilder.helper import LcbuilderHelper
from lcbuilder.lcbuilder_class import LcBuilder
from lcbuilder.objectinfo.MissionFfiIdObjectInfo import MissionFfiIdObjectInfo
from lcbuilder.objectinfo.MissionObjectInfo import MissionObjectInfo
from lcbuilder.photometry.aperture_extractor import ApertureExtractor
from lcbuilder.star.HabitabilityCalculator import HabitabilityCalculator
from lightkurve import TessLightCurve, TessTargetPixelFile, KeplerTargetPixelFile
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 scipy import stats, ndimage

from watson import constants
import watson.tpfplotterSub.tpfplotter as tpfplotter
import pandas as pd
import os
from math import ceil

from watson.report import Report

class Watson:
    """
    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):
        self.object_dir = os.getcwd() if object_dir is None else object_dir
        self.data_dir = self.object_dir + "/"

    def vetting(self, id, period, t0, duration, depth, ffi, sectors, rp_rstar=None, a_rstar=None, cpus=None,
                cadence=None, lc_file=None, lc_data_file=None, tpfs_dir=None, apertures_file=None,
                create_fov_plots=False, cadence_fov=None, ra=None, dec=None, transits_list=None,
                v=None, j=None, h=None, k=None, clean=True, transits_mask=None):
        """
        Launches the whole vetting procedure that ends up with a validation report
        :param id: the target star id
        :param period: the period of the candidate in days
        :param t0: the epoch in days
        :param duration: the duration of the transit of the candidate in minutes
        :param depth: the depth of the transit of the candidate in ppts
        :param ffi: flag to specify whether the data to be used is FFI or short cadence curve
        :param sectors: sectors/quarters/campaigns to be used
        :param rp_rstar: Rp / Rstar
        :param a_rstar: Semi-major axis / Rstar
        :param cpus: number of cpus to be used
        :param cadence: the cadence to be used to download data, in seconds
        :param lc_file: the file containing the curve
        :param lc_data_file: the file containing the raw curve and the motion, centroids and quality flags
        :param tpfs_dir: the directory containing the tpf files
        :param apertures_file: the file containing the map of sectors->apertures
        :param create_fov_plots: whether to generate Field Of View plots.
        :param cadence_fov: the cadence to use to download fov_plots
        :param ra: the RA to use to download fov_plots
        :param dec: the DEC to use to download fov_plots
        :param transits_list: a list of dictionaries with shape: {'t0': value, 'depth': value, 'depth_err': value}
        :param v: star V magnitude
        :param j: star J magnitude
        :param h: star H magnitude
        :param k: star K magnitude
        :param clean: whether to clean all the pngs created for the final pdfs
        :param transits_mask: array with shape [{P:period, T0:t0, D:d}, ...] to use for transits masking before vetting
        """
        logging.info("------------------")
        logging.info("Candidate info")
        logging.info("------------------")
        logging.info("Id: %.s", id)
        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("Rp_Rstar: %.4f", rp_rstar)
        logging.info("a_Rstar: %.2f", a_rstar)
        logging.info("FFI: %s", ffi)
        logging.info("Sectors: %s", sectors)
        lc_builder = LcBuilder()
        if transits_mask is None:
            transits_mask = []
        if rp_rstar is None:
            rp_rstar = np.sqrt(depth / 1000)
        lc_build = None
        if lc_file is None or lc_data_file is None:
            if not ffi:
                lc_build = lc_builder.build(MissionObjectInfo(id, sectors, cadence=cadence,
                                                              initial_transit_mask=transits_mask), self.data_dir)
            else:
                lc_build = lc_builder.build(MissionFfiIdObjectInfo(id, sectors, cadence=cadence,
                                                              initial_transit_mask=transits_mask), self.data_dir)
            lc_build.lc_data.to_csv(self.data_dir + "/lc_data.csv")
            lc_file = self.data_dir + "/lc.csv"
            lc_data_file = self.data_dir + "/lc_data.csv"
        if a_rstar is None and lc_build is None:
            raise ValueError("You need to define a_rstar if you are providing the lc_file and lc_data_file")
        if a_rstar is None:
            a_rstar = HabitabilityCalculator().calculate_semi_major_axis(period, lc_build.star_info.mass)
        if tpfs_dir is None:
            tpfs_dir = self.data_dir + "/tpfs/"
        if apertures_file is None:
            apertures_file = self.data_dir + "/apertures.yaml"
        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
        try:
            transits_list_t0s, summary_list_t0s_indexes = self.__process(id, period, t0, duration, depth, rp_rstar, a_rstar,
                                                                 cpus, lc_file, lc_data_file, tpfs_dir,
                                                                 apertures_file, create_fov_plots, cadence_fov, ra,
                                                                 dec, transits_list, transits_mask)
            self.report(id, ra, dec, t0, period, duration, depth, transits_list_t0s, summary_list_t0s_indexes,
                        v, j, h, k, os.path.exists(tpfs_dir))
            if clean:
                for filename in os.listdir(self.data_dir):
                    if not filename.endswith(".pdf"):
                        os.remove(self.data_dir + "/" + filename)

        except Exception as e:
            traceback.print_exc()

    def report(self, id, ra, dec, t0, period, duration, depth, transits_list, summary_list_t0s_indexes, v, j, h, k,
               with_tpfs=True):
        file_name = "transits_validation_report.pdf"
        report = Report(self.data_dir, file_name, id, ra, dec, t0, period, duration, depth, transits_list, None,
                        v, j, h, k, with_tpfs)
        report.create_report()
        file_name = "transits_validation_report_summary.pdf"
        report = Report(self.data_dir, file_name, id, ra, dec, t0, period, duration, depth, transits_list,
                        summary_list_t0s_indexes, v, j, h, k, with_tpfs)
        report.create_report()

    def vetting_with_data(self, candidate_df, star, transits_df, cpus, create_fov_plots=False, cadence_fov=None,
                          transits_mask=None):
        """
        Same than vetting but receiving a candidate dataframe and a star dataframe with one row each.
        :param candidate_df: the candidate dataframe containing id, period, t0, transits and sectors data.
        :param star: the star dataframe with the star info.
        :param transits_df: a dataframe containing the transits information with columns 't0', 'depth' and 'depth_err'
        :param cpus: the number of cpus to be used.
        :param create_fov_plots: whether to generate Field Of View plots.
        :param cadence_fov: the cadence to use to download fov_plots
        :param transits_mask: array with shape [{P:period, T0:t0, D:d}, ...] to use for transits masking before vetting
        """
        if transits_mask is None:
            transits_mask = []
        df = candidate_df.iloc[0]
        # TODO get the transit time list
        id = df['id']
        period = df['period']
        t0 = df['t0']
        rp_rstar = df['rp_rs']
        a_rstar = df['a'] / star["R_star"] * constants.AU_TO_RSUN
        duration = df['duration']
        depth = df['depth']
        ffi = df['ffi']
        run = int(df['number'])
        curve = int(df['curve'])
        sectors = df['sectors']
        if isinstance(sectors, (int, np.integer)):
            sectors = [sectors]
        elif isinstance(sectors, (str)):
            sectors = sectors.split(',')
        lc_file = "/lc_" + str(curve) + ".csv"
        lc_file = self.object_dir + lc_file
        lc_data_file = self.object_dir + "/lc_data.csv"
        tpfs_dir = self.object_dir + "/tpfs"
        apertures_file = self.object_dir + "/apertures.yaml"
        try:
            self.vetting(id, period, t0, duration, depth, ffi, sectors, rp_rstar=rp_rstar, a_rstar=a_rstar, cpus=cpus,
                         lc_file=lc_file, lc_data_file=lc_data_file, tpfs_dir=tpfs_dir, apertures_file=apertures_file,
                         create_fov_plots=create_fov_plots, cadence_fov=cadence_fov, ra=star["ra"],
                         dec=star["dec"], transits_list=transits_df.to_dict("list"), transits_mask=transits_mask)
        except Exception as e:
            traceback.print_exc()

    def __process(self, id, period, t0, duration, depth, rp_rstar, a_rstar, cpus, lc_file, lc_data_file, tpfs_dir,
                  apertures_file, create_fov_plots=False, cadence_fov=None, ra_fov=None, dec_fov=None,
                  transits_list=None, transits_mask=None):
        """
        Performs the analysis to generate PNGs and Transits Validation Report.
        :param id: the target star id
        :param period: the period of the candidate in days
        :param t0: the epoch in days
        :param duration: the duration of the transit of the candidate in minutes
        :param depth: the depth of the transit of the candidate in ppts
        :param sectors: sectors/quarters/campaigns to be used
        :param rp_rstar: Rp / Rstar
        :param a_rstar: Semi-major axis / Rstar
        :param cpus: number of cpus to be used
        :param lc_file: the file containing the curve
        :param lc_data_file: the file containing the raw curve and the motion, centroids and quality flags
        :param tpfs_dir: the directory containing the tpf files
        :param apertures_file: the file containing the apertures
        :param create_fov_plots: whether to create FOV plots
        :param cadence_fov: the cadence to use to download fov_plots
        :param ra_fov: the RA to use to download fov_plots
        :param dec_fov: the DEC to use to download fov_plots
        :param transits_list: a list of dictionaries with shape: {'t0': value, 'depth': value, 'depth_err': value}
        :param transits_mask: array with shape [{P:period, T0:t0, D:d}, ...] to use for transits masking before vetting
        """
        logging.info("Running Transit Plots")
        lc, lc_data, tpfs = Watson.initialize_lc_and_tpfs(id, lc_file, lc_data_file, tpfs_dir,
                                                          transits_mask=transits_mask)
        apertures = None
        if os.path.exists(apertures_file):
            apertures = yaml.load(open(apertures_file), yaml.SafeLoader)
            apertures = apertures["sectors"]
            mission, mission_prefix, mission_int_id = LcBuilder().parse_object_info(id)
            if create_fov_plots:
                if cadence_fov is None:
                    cadence_fov = LcbuilderHelper.compute_cadence(lc.time.value)
                Watson.vetting_field_of_view(self.data_dir, mission, mission_int_id, cadence_fov, ra_fov, dec_fov,
                                             list(apertures.keys()), "tpf", apertures, cpus)
        summary_t0s_indexes = None
        if transits_list is not None:
            transits_list_not_nan_indexes = \
                Watson.plot_transits_statistics(self.data_dir, id, t0, period, transits_list)
            transit_t0s_list = np.array(transits_list["t0"])[transits_list_not_nan_indexes]
            transit_depths = np.array(transits_list["depth"])[transits_list_not_nan_indexes]
            summary_t0s_indexes = np.argwhere((transit_depths == np.max(transits_list["depth"])) |
                                              (transit_depths == np.min(transits_list["depth"]))).flatten()
            if len(transit_depths) > 2:
                closest_depths_to_mean = np.abs(transit_depths - depth)
                summary_t0s_indexes = np.append(summary_t0s_indexes, np.argmin(closest_depths_to_mean))
        else:
            last_time = lc.time.value[len(lc.time.value) - 1]
            num_of_transits = int(ceil(((last_time - t0) / period)))
            transit_lists = t0 + period * np.arange(0, num_of_transits)
            time_as_array = lc.time.value
            plot_range = duration / 3600 * 2
            transits_in_data = [time_as_array[(transit > time_as_array - plot_range) & (transit < time_as_array + plot_range)] for
                                transit in transit_lists]
            transit_t0s_list = transit_lists[[len(transits_in_data_set) > 0 for transits_in_data_set in transits_in_data]]
        self.plot_folded_curve(self.data_dir, id, lc, period, t0, duration, depth / 1000, rp_rstar, a_rstar)
        plot_transits_inputs = []
        for index, transit_times in enumerate(transit_t0s_list):
            plot_transits_inputs.append(SingleTransitProcessInput(self.data_dir, str(id), index, lc_file, lc_data_file,
                                                                  tpfs_dir, apertures, transit_times, depth / 1000,
                                                                  duration, period, rp_rstar, a_rstar))
        with multiprocessing.Pool(processes=cpus) as pool:
            pool.map(Watson.plot_single_transit, plot_transits_inputs)
        return transit_t0s_list, summary_t0s_indexes

    @staticmethod
    def initialize_lc_and_tpfs(id, lc_file, lc_data_file, tpfs_dir, transits_mask=None):
        if transits_mask is None:
            transits_mask = []
        mission, mission_prefix, mission_int_id = LcBuilder().parse_object_info(id)
        lc = pd.read_csv(lc_file, header=0)
        lc_data = None
        if os.path.exists(lc_data_file):
            lc_data = pd.read_csv(lc_data_file, header=0)
            lc_data = Watson.normalize_lc_data(lc_data)
        time, flux, flux_err = lc["#time"].values, lc["flux"].values, lc["flux_err"].values
        for transit_mask in transits_mask:
            logging.info('* Transit mask with P=%.2f d, T0=%.2f d, Dur=%.2f min *', transit_mask["P"],
                         transit_mask["T0"], transit_mask["D"])
            mask = foldedleastsquares.transit_mask(time, transit_mask["P"], transit_mask["D"] / 60 / 24,
                                                   transit_mask["T0"])
            time = time[~mask]
            flux = flux[~mask]
            flux_err = flux_err[~mask]
        lc = TessLightCurve(time=time, flux=flux, flux_err=flux_err, quality=np.zeros(len(time)))
        lc.extra_columns = []
        tpfs = []
        if os.path.exists(tpfs_dir):
            for tpf_file in os.listdir(tpfs_dir):
                tpf = TessTargetPixelFile(tpfs_dir + "/" + tpf_file) if mission == lcbuilder.constants.MISSION_TESS else \
                    KeplerTargetPixelFile(tpfs_dir + "/" + tpf_file)
                tpfs.append(tpf)
        return lc, lc_data, tpfs

    @staticmethod
    def normalize_lc_data(lc_data):
        logging.info("Normalizing lc_data")
        time = lc_data["time"].to_numpy()
        dif = time[1:] - time[:-1]
        jumps = np.where(np.abs(dif) > 1)[0]
        jumps = np.append(jumps, len(lc_data))
        previous_jump_index = 0
        for jumpIndex in jumps:
            token = lc_data["centroids_x"][previous_jump_index:jumpIndex]
            lc_data["centroids_x"][previous_jump_index:jumpIndex] = token - np.nanmedian(token)
            token = lc_data["centroids_y"][previous_jump_index:jumpIndex]
            lc_data["centroids_y"][previous_jump_index:jumpIndex] = token - np.nanmedian(token)
            token = lc_data["motion_x"][previous_jump_index:jumpIndex]
            lc_data["motion_x"][previous_jump_index:jumpIndex] = token - np.nanmedian(token)
            token = lc_data["motion_y"][previous_jump_index:jumpIndex]
            lc_data["motion_y"][previous_jump_index:jumpIndex] = token - np.nanmedian(token)
            previous_jump_index = jumpIndex
        return lc_data

    @staticmethod
    def plot_transits_statistics(data_dir, id, epoch, period, transits_list):
        fig, axs = plt.subplots(1, 1, figsize=(12, 6), constrained_layout=True)
        fig.suptitle(str(id) + ' Transits depth analysis T0=' + str(round(epoch, 2)) + ' P=' + str(round(period, 2)) + 'd', size=18)
        transits_list_not_nan_t0s_indexes = np.argwhere(~np.isnan(transits_list["t0"])).flatten()
        transits_list_not_nan_depths_indexes = np.argwhere(~np.isnan(transits_list["depth"])).flatten()
        transits_list_not_nan_indexes = np.intersect1d(transits_list_not_nan_t0s_indexes, transits_list_not_nan_depths_indexes)
        transits_list_t0s = np.array(transits_list["t0"])[transits_list_not_nan_indexes]
        transits_list_depths = np.array(transits_list["depth"])[transits_list_not_nan_indexes]
        transits_list_depths_err = np.array(transits_list["depth_err"])[transits_list_not_nan_indexes]
        even_transits_indexes = np.argwhere((np.abs((transits_list_t0s - epoch) % (2 * period)) < 0.1) |
                                            (np.abs((transits_list_t0s - epoch) % (2 * period)) > (
                                                        2 * period) - 0.1)).flatten()
        odd_transits_indexes = np.argwhere((np.abs((transits_list_t0s - epoch) % (2 * period)) > period - 0.05) &
                                           (np.abs((transits_list_t0s - epoch) % (2 * period)) < period + 0.05)).flatten()
        axs.axhline(y=np.mean(transits_list_depths), color='purple', alpha=0.3,
                    ls='-', lw=2, label='Depth Mean')
        axs.axhline(y=np.percentile(transits_list_depths, 84), color='purple', alpha=0.3,
                    ls='--', lw=2, label='Depth 1-sigma confidence')
        axs.axhline(y=np.percentile(transits_list_depths, 18), color='purple', alpha=0.3,
                    ls='--', lw=2)
        axs.axhline(y=np.mean(transits_list_depths[even_transits_indexes]), color='blue', alpha=0.3,
                    ls='-', lw=2, label='Depth Mean Even')
        axs.axhline(y=np.mean(transits_list_depths[odd_transits_indexes]), color='red', alpha=0.3,
                    ls='-', lw=2, label='Depth Mean Odd')
        axs.errorbar(x=even_transits_indexes, y=transits_list_depths[even_transits_indexes],
                     yerr=transits_list_depths_err[even_transits_indexes],
                     fmt="o", color="blue", ecolor="cyan", label="Even transits")
        axs.errorbar(x=odd_transits_indexes, y=transits_list_depths[odd_transits_indexes],
                     yerr=transits_list_depths_err[odd_transits_indexes],
                     fmt="o", color="red", ecolor="darkorange", label="Odd transits")
        axs.legend(loc='center left', bbox_to_anchor=(1, 0.5))
        axs.set_xlabel("Transit number")
        axs.set_ylabel("Depth (ppt)")
        plt.savefig(data_dir + "/transit_depths.png")
        plt.clf()
        plt.close()
        return transits_list_not_nan_indexes

    @staticmethod
    def plot_single_transit(single_transit_process_input):
        """
        Plots the single transit info: single transit focused curve, drift and background plots, small vs used aperture
        photometry and tpf flux values around the transit.
        @param single_transit_process_input: wrapper class to provide pickable inputs for multiprocessing
        """
        lc, lc_data, tpfs = Watson.initialize_lc_and_tpfs(single_transit_process_input.id,
                                                          single_transit_process_input.lc_file,
                                                          single_transit_process_input.lc_data_file,
                                                          single_transit_process_input.tpfs_dir)
        transit_time = single_transit_process_input.transit_times
        duration = single_transit_process_input.duration / 60 / 24
        fig = plt.figure(figsize=(24, 6), constrained_layout=True)
        fig.suptitle('Vetting of ' + str(single_transit_process_input.id) + ' single transit no. ' +
                     str(single_transit_process_input.index) +
                     ' at T0=' + str(round(transit_time, 2)) + 'd', size=26)
        gs = gridspec.GridSpec(2, 3, hspace=0.4, wspace=0.1)  # 2 rows, 3 columns
        ax1 = fig.add_subplot(gs[0, 0])  # First row, first column
        ax2 = fig.add_subplot(gs[0, 1])  # First row, second column
        ax3 = fig.add_subplot(gs[0, 2])  # First row, third column
        ax4 = fig.add_subplot(gs[1, 0])  # First row, third column
        ax5 = fig.add_subplot(gs[1, 1])  # First row, third column
        ax6 = fig.add_subplot(gs[1, 2])  # First row, third column
        axs = [ax1, ax2, ax3, ax4, ax5, ax6]
        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]
        plot_range = duration * 2
        zoom_mask = np.where((time > transit_time - plot_range) & (time < transit_time + plot_range))
        plot_time = time[zoom_mask]
        plot_flux = flux[zoom_mask]
        zoom_lc_data = None
        if lc_data is not None:
            zoom_lc_data = lc_data[(lc_data["time"] > transit_time - plot_range) & (lc_data["time"] < transit_time +
                                                                                plot_range)]
        aperture_mask = None
        chosen_aperture_lc = None
        smaller_aperture_lc = None
        eroded_aperture_mask = None
        tpf_short_framed = None
        if len(plot_time) > 0:
            for tpf in tpfs:
                tpf_zoom_mask = tpf[(tpf.time.value > transit_time - plot_range) & (tpf.time.value < transit_time +
                                                                                plot_range)]
                if len(tpf_zoom_mask) > 0:
                    if hasattr(tpf, 'sector') and tpf.sector is not None:
                        sector = tpf.sector
                    elif hasattr(tpf, 'campaign') and tpf.campaign:
                        sector = tpf.campaign
                    else:
                        sector = tpf.quarter
                    tpf_short_framed = tpf[(tpf.time.value > transit_time - plot_range) &
                                           (tpf.time.value < transit_time + plot_range)]
                    if len(tpf_short_framed) == 0:
                        break
                    if not isinstance(single_transit_process_input.apertures, (dict)) and \
                            np.isnan(single_transit_process_input.apertures):
                        chosen_aperture_lc = lc
                        smaller_aperture_lc = lc
                        aperture_mask = [[]]
                    else:
                        aperture_mask = ApertureExtractor.from_pixels_to_boolean_mask(
                            single_transit_process_input.apertures[sector], tpf.column, tpf.row, tpf.shape[2], tpf.shape[1])
                        eroded_aperture_mask = ndimage.binary_erosion(aperture_mask)
                        chosen_aperture_lc = tpf.to_lightcurve(aperture_mask=aperture_mask)
                    if True in eroded_aperture_mask:
                        smaller_aperture_lc = tpf.to_lightcurve(aperture_mask=eroded_aperture_mask)
                    break
            single_transit_file = single_transit_process_input.data_dir + "/single_transit_" + \
                                  str(single_transit_process_input.index) + "_T0_" + str(transit_time) + ".png"
            tpf_single_transit_file = single_transit_process_input.data_dir + "/tpf_single_transit_" + \
                                      str(single_transit_process_input.index) + "_T0_" + str(transit_time) + ".png"
            t1 = transit_time - duration / 2
            t4 = transit_time + duration / 2
            model_time, model_flux = Watson.get_transit_model(duration, transit_time,
                                                              (transit_time - plot_range,
                                                               transit_time + plot_range),
                                                              single_transit_process_input.depth,
                                                              single_transit_process_input.period,
                                                              single_transit_process_input.rp_rstar,
                                                              single_transit_process_input.a_rstar)
            momentum_dumps_lc_data = None
            if zoom_lc_data is not None and not zoom_lc_data["quality"].isnull().all():
                momentum_dumps_lc_data = zoom_lc_data[np.bitwise_and(zoom_lc_data["quality"].to_numpy(),
                                                                     constants.MOMENTUM_DUMP_QUALITY_FLAG) >= 1]

            axs[0].plot(model_time, model_flux, color="red")
            axs[0].scatter(plot_time, plot_flux, color="darkorange", label="Photometry used aperture")
            axs[0].set_xlim([transit_time - plot_range, transit_time + plot_range])
            axs[0].set_title("Single Transit")
            if momentum_dumps_lc_data is not None and len(momentum_dumps_lc_data) > 0:
                first_momentum_dump = True
                for index, momentum_dump_row in momentum_dumps_lc_data.iterrows():
                    if first_momentum_dump:
                        axs[0].axvline(x=momentum_dump_row["time"], color='purple', alpha=0.3,
                                       ls='--', lw=2, label='Momentum dump')
                        first_momentum_dump = False
                    else:
                        axs[0].axvline(x=momentum_dump_row["time"], color='purple', alpha=0.3,
                                       ls='--', lw=2)
            axs[0].legend(loc='upper left')
            axs[0].set_xlim([transit_time - plot_range, transit_time + plot_range])
            axs[0].set_xlabel("Time (d)")
            axs[0].set_ylabel("Flux norm.")
            if zoom_lc_data is not None:
                axs[1].scatter(zoom_lc_data["time"], zoom_lc_data["motion_x"], color="red",
                               label="X-axis motion")
                axs[1].scatter(zoom_lc_data["time"], zoom_lc_data["centroids_x"],
                               color="black", label="X-axis centroids")
                axs[1].axvline(x=transit_time - duration / 2, color='r', label='T1')
                axs[1].axvline(x=transit_time + duration / 2, color='r', label='T4')
                axs[1].legend(loc='upper left')
            axs[1].set_xlim([transit_time - plot_range, transit_time + plot_range])
            axs[1].set_title("X-axis drift")
            axs[1].set_xlabel("Time (d)")
            axs[1].set_ylabel("Normalized Y-axis data")
            if zoom_lc_data is not None:
                axs[2].scatter(zoom_lc_data["time"], zoom_lc_data["motion_y"], color="red",
                               label="Y-axis motion")
                axs[2].scatter(zoom_lc_data["time"], zoom_lc_data["centroids_y"],
                               color="black", label="Y-axis centroids")
                axs[2].axvline(x=t1, color='r', label='T1')
                axs[2].axvline(x=t4, color='r', label='T4')
                axs[2].legend(loc='upper left')
            axs[2].set_xlim([transit_time - plot_range, transit_time + plot_range])
            axs[2].set_title("Y-axis drift")
            axs[2].set_xlabel("Time (d)")
            axs[2].set_ylabel("Normalized Y-axis data")
            if smaller_aperture_lc is not None:
                axs[3].plot(model_time, model_flux, color="red")
                axs[3].set_xlim([transit_time - plot_range, transit_time + plot_range])
                axs[3].set_title("SAP comparison")
                axs[3].set_xlim([transit_time - plot_range, transit_time + plot_range])
                axs[3].set_xlabel("Time (d)")
                axs[3].set_ylabel("Flux norm.")
                chosen_aperture_lc.flux = wotan.flatten(chosen_aperture_lc.time.value,
                                                        chosen_aperture_lc.flux.value, window_length=0.75,
                                                        return_trend=False, method="biweight", break_tolerance=0.5)
                smaller_aperture_lc.flux = wotan.flatten(smaller_aperture_lc.time.value,
                                                         smaller_aperture_lc.flux.value, window_length=0.75,
                                                         return_trend=False, method="biweight", break_tolerance=0.5)
                chosen_aperture_lc = chosen_aperture_lc[
                    (chosen_aperture_lc.time.value - plot_range < transit_time) &
                    (chosen_aperture_lc.time.value + plot_range > transit_time)]
                smaller_aperture_lc = smaller_aperture_lc[
                    (smaller_aperture_lc.time.value - plot_range < transit_time) &
                    (smaller_aperture_lc.time.value + plot_range > transit_time)]
                axs[3].scatter(chosen_aperture_lc.time.value, chosen_aperture_lc.flux.value, color="darkorange",
                               label="Photometry used aperture")
                axs[3].scatter(smaller_aperture_lc.time.value, smaller_aperture_lc.flux.value,
                               color="c", label="Photometry smaller aperture")
            axs[3].legend(loc='upper left')
            if tpf_short_framed is not None:
                axs[4] = tpf_short_framed.plot(axs[4], aperture_mask=aperture_mask)
                axs[4].set_title("TPF apertures comparison")
                if smaller_aperture_lc is not None:
                    parsed_aperture = tpf_short_framed._parse_aperture_mask(eroded_aperture_mask)
                    for i in range(tpf_short_framed.shape[1]):
                        for j in range(tpf_short_framed.shape[2]):
                            if parsed_aperture[i, j]:
                                rect = patches.Rectangle(
                                    xy=(j + tpf_short_framed.column - 0.5, i + tpf_short_framed.row - 0.5),
                                    width=1,
                                    height=1,
                                    color='black',
                                    fill=False,
                                    hatch="\\\\",
                                )
                                axs[4].add_patch(rect)
            if zoom_lc_data is not None:
                axs[5].scatter(zoom_lc_data["time"], zoom_lc_data["background_flux"],
                               color="blue", label="Background Flux")
                axs[5].axvline(x=t1, color='r', label='T1')
                axs[5].axvline(x=t4, color='r', label='T4')
                axs[5].legend(loc='upper left')
            axs[5].set_xlim([transit_time - plot_range, transit_time + plot_range])
            axs[5].set_title("Background flux")
            axs[5].set_xlabel("Time (d)")
            axs[5].set_ylabel("Background flux (e/s)")
            plt.savefig(single_transit_file, dpi=100, bbox_inches='tight')
            plt.clf()
            plt.close()
            if tpf_short_framed is not None:
                tpf_short_framed.plot_pixels(aperture_mask=aperture_mask)
                plt.savefig(tpf_single_transit_file, dpi=100)
                plt.clf()
                plt.close()
                images_list = [single_transit_file, tpf_single_transit_file]
                imgs = [PIL.Image.open(i) for i in images_list]
                imgs[0] = imgs[0].resize((imgs[1].size[0],
                                          int(imgs[1].size[0] / imgs[0].size[0] * imgs[0].size[1])),
                                          PIL.Image.ANTIALIAS)
                img_merge = np.vstack((np.asarray(i) for i in imgs))
                img_merge = PIL.Image.fromarray(img_merge)
                img_merge.save(single_transit_file, quality=95, optimize=True)
                os.remove(tpf_single_transit_file)
            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, rp_rstar, a_rstar):
        """
        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
        Watson.compute_phased_values_and_fill_plot(id, axs[0][0], lc, period, epoch + period / 2, depth, duration,
                                                   rp_rstar, a_rstar, bins=bins)
        Watson.compute_phased_values_and_fill_plot(id, axs[0][1], lc, period, epoch, depth, duration, rp_rstar, a_rstar, bins=bins)
        period = 2 * period
        Watson.compute_phased_values_and_fill_plot(id, axs[1][0], lc, period, epoch + period / 2, depth, duration,
                                                   rp_rstar, a_rstar, bins=bins)
        Watson.compute_phased_values_and_fill_plot(id, axs[1][1], lc, period, epoch, depth, duration, rp_rstar, a_rstar, bins=bins)
        period = period / 4
        Watson.compute_phased_values_and_fill_plot(id, axs[2][0], lc, period, epoch + period / 2, depth, duration,
                                                   rp_rstar, a_rstar, bins=bins)
        Watson.compute_phased_values_and_fill_plot(id, axs[2][1], lc, period, epoch, depth, duration, rp_rstar, a_rstar, 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, rp_rstar, a_rstar, 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 and len(folded_y) > bins:
            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)
            bin_nan_args = np.isnan(bin_stds)
            axs.errorbar(bin_centers[~bin_nan_args], bin_means[~bin_nan_args],
                         yerr=bin_stds[~bin_nan_args] / 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
        model_time, model_flux = Watson.get_transit_model(half_duration_phase * 2, 0.5,
                                                          (0.5 - half_duration_phase * range, 0.5 + half_duration_phase * range),
                                                          depth, period, rp_rstar, a_rstar, 2 * len(time))
        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, period, rp_to_rstar, a_to_rstar, 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 = rp_to_rstar  # planet radius (in units of stellar radii)
        ma.a = a_to_rstar  # 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_intransit = np.argwhere(model < 1)[:, 0]
        model_time = np.linspace(start_end[0], start_end[1], len(model))
        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

    @staticmethod
    def plot_tpf(tpf, sector, aperture, dir):
        logging.info("Plotting FOV curves for sector %.0f", sector)
        if not os.path.exists(dir):
            os.mkdir(dir)
        tpf.plot_pixels(aperture_mask=aperture)
        plt.savefig(dir + "/fov_Flux_pixels[" + str(sector) + "].png")
        plt.close()

    @staticmethod
    def compute_pixels_curves(tpf):
        masks = np.zeros(
            (tpf.shape[1] * tpf.shape[2], tpf.shape[1], tpf.shape[2]),
            dtype="bool",
        )
        for i in range(tpf.shape[1] * tpf.shape[2]):
            masks[i][np.unravel_index(i, (tpf.shape[1], tpf.shape[2]))] = True
        pixel_list = []
        for j in range(tpf.shape[1] * tpf.shape[2]):
            lc = tpf.to_lightcurve(aperture_mask=masks[j])
            lc = lc.remove_outliers(sigma_upper=3, sigma_lower=float('inf'))
            if len(lc.remove_nans().flux) == 0:
                pixel_list.append(None)
            else:
                pixel_list.append(lc)

    @staticmethod
    def vetting_field_of_view_single(fov_process_input):
        """
        Plots FOV for one sector data. To be called by a multiprocessing queue.
        :param fov_process_input: wrapper for the sector data
        """
        maglim = 6
        search_radius = 40
        try:
            tpf = fov_process_input.tpf_source.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 = fov_process_input.apertures[tpf.sector]
            aperture = aperture if isinstance(aperture, np.ndarray) else np.array(aperture)
            aperture_boolean = ApertureExtractor.from_pixels_to_boolean_mask(aperture, column, row, tpf.shape[2],
                                                                             tpf.shape[1])
            Watson.plot_tpf(tpf, tpf.sector, aperture_boolean, fov_process_input.save_dir)
            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(fov_process_input.ra, fov_process_input.dec,
                                                    search_radius=search_radius)
            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 ' + fov_process_input.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)
            plt.savefig(fov_process_input.save_dir + '/fov_TPF_Gaia_' + fov_process_input.target_title + '_S' +
                        str(tpf.sector) + '.png')
            # 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, fov_process_input.save_dir + '/fov_Gaia_' + fov_process_input.target_title + '_S' +
                        str(tpf.sector) + '.dat', overwrite=True)
        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")

    @staticmethod
    def vetting_field_of_view(indir, mission, tic, cadence, ra, dec, sectors, source, apertures,
                              cpus=multiprocessing.cpu_count() - 1):
        """
        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
        :param cpus: cores to be used
        :return: the directory where resulting data is stored
        """
        try:
            sectors = [sectors] if isinstance(sectors, int) else sectors
            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")
            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:
                    ra_str = str(ra)
                    dec_str = "+" + str(dec) if dec >= 0 else str(dec)
                    coords_str = ra_str + " " + dec_str
                    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
            if not os.path.exists(save_dir):
                os.makedirs(save_dir)
            fov_process_inputs = []
            for i in range(0, len(tpf_source)):
                fov_process_inputs.append(FovProcessInput(save_dir, mission, tic, cadence, ra, dec, sectors, source,
                                                          apertures, tpf_source[i], target_title))
            with multiprocessing.Pool(processes=cpus) as pool:
                pool.map(Watson.vetting_field_of_view_single, fov_process_inputs)
            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")

class SingleTransitProcessInput:
    def __init__(self, data_dir, id, index, lc_file, lc_data_file, tpfs_dir, apertures,
                                         transit_times, depth, duration, period, rp_rstar, a_rstar):
        self.data_dir = data_dir
        self.id = id
        self.index = index
        self.lc_file = lc_file
        self.lc_data_file = lc_data_file
        self.tpfs_dir = tpfs_dir
        self.apertures = apertures
        self.transit_times = transit_times
        self.depth = depth
        self.duration = duration
        self.period = period
        self.rp_rstar = rp_rstar
        self.a_rstar = a_rstar

class FovProcessInput:
    def __init__(self, save_dir, mission, tic, cadence, ra, dec, sectors, source, apertures, tpf_source, target_title):
        self.save_dir = save_dir
        self.mission = mission
        self.tic = tic
        self.cadence = cadence
        self.ra = ra
        self.dec = dec
        self.sectors = sectors
        self.source = source
        self.apertures = apertures
        self.tpf_source = tpf_source
        self.target_title = target_title

import numpy,from scipy,import astropy,from astropy


import os
import itertools
from copy import deepcopy

import numpy as np
import torch

from .nn_blocks import Actor, Critic, make_feature_extractor_fn
from .noise import GaussianNoise, OUNoise

def soft_update_net(args, online_net, target_net):
    tau = args.tau
    with torch.no_grad():
        for online_x, target_x in zip(online_net.parameters(), target_net.parameters()):
            target_x.data.mul_(1 - tau)
            target_x.data.add_(tau * online_x.data)

class TD3Agent(object):
    def __init__(self, args, env, device):
        self.device = device
        feature_extractor_fn = make_feature_extractor_fn(args, env)
        self.online_actor = Actor(args, env.n_act, feature_extractor_fn).to(device)
        self.actor_optimizer = torch.optim.Adam(self.online_actor.parameters(), lr=args.actor_lr)

        self.online_Q1, self.online_Q2 = Critic(args, env.n_act, feature_extractor_fn).to(device), Critic(args, env.n_act, feature_extractor_fn).to(device)
        self.Q_params = itertools.chain(self.online_Q1.parameters(), self.online_Q2.parameters())
        self.Q_optimizer = torch.optim.Adam(self.Q_params, lr=args.critic_lr)

        self.target_actor, self.target_Q1, self.target_Q2 = deepcopy(self.online_actor), deepcopy(self.online_Q1), deepcopy(self.online_Q2)
        for x in self.target_actor.parameters():
            x.requires_grad = False
        for x in self.target_Q1.parameters():
            x.requires_grad = False
        for x in self.target_Q2.parameters():
            x.requires_grad = False

        self.learn_steps = 0
        noise_gen_seed = np.random.randint(1e9)
        self.noise_gen = OUNoise(env.n_act, noise_gen_seed) if args.OU_noise else GaussianNoise(env.n_act, noise_gen_seed)

    @torch.no_grad()
    def act(self, obs, noise_scale=0):
        obs = torch.tensor(obs, dtype=torch.float32, device=self.device).unsqueeze(0)
        action = self.online_actor(obs).squeeze(0).cpu().numpy()
        action += noise_scale * self.noise_gen()
        return np.clip(action, -1, 1)

    def lr_decay(self, args, cur, total):
        for param_group in self.actor_optimizer.param_groups:
            param_group["lr"] = args.actor_lr * (1 - cur / total)
        for param_group in self.Q_optimizer.param_groups:
            param_group["lr"] = args.critic_lr * (1 - cur / total)

    def learn(self, args, buffer, T):
        stats = {}
        batch_size = args.batch_size

        self.learn_steps += 1
        idx = np.random.choice(len(buffer), batch_size, replace=True)
        data = buffer.get(idx, collect_next_obs=True)
        obs = torch.as_tensor(data["obs"]).to(self.device)
        action = torch.as_tensor(data["action"]).to(self.device)
        reward = torch.as_tensor(data["reward"]).to(self.device)
        done = torch.as_tensor(data["done"]).to(self.device)
        next_obs = torch.as_tensor(data["next_obs"]).to(self.device)

        Q1, Q2 = self.online_Q1(obs, action), self.online_Q2(obs, action)
        with torch.no_grad():
            next_action = self.target_actor(next_obs)
            next_eps = (args.target_noise_scale * torch.randn_like(next_action)).clamp(-args.target_noise_clip, args.target_noise_clip)
            next_action = (next_action + next_eps).clamp(-1, 1)

            next_Q1, next_Q2 = self.target_Q1(next_obs, next_action), self.target_Q2(next_obs, next_action)
            next_Q = torch.min(next_Q1, next_Q2)
            tar_Q = reward + args.gamma * (1 - done) * next_Q

        Q_loss = ((Q1 - tar_Q) ** 2).mean() + ((Q2 - tar_Q) ** 2).mean()
        self.Q_optimizer.zero_grad()
        Q_loss.backward()
        self.Q_optimizer.step()
        stats["Q1"] = Q1.mean().item()
        stats["Q2"] = Q2.mean().item()
        stats["Q_loss"] = Q_loss.item()

        if self.learn_steps % args.update_delay == 0:
            for x in self.Q_params:
                x.requires_grad = False

            actor_loss = -self.online_Q1(obs, self.online_actor(obs)).mean()
            self.actor_optimizer.zero_grad()
            actor_loss.backward()
            self.actor_optimizer.step()
            stats["actor_loss"] = actor_loss.item()

            for x in self.Q_params:
                x.requires_grad = True
            soft_update_net(args, self.online_actor, self.target_actor)
            soft_update_net(args, self.online_Q1, self.target_Q1)
            soft_update_net(args, self.online_Q2, self.target_Q2)

        if args.wandb_show:
            from .logger import log
            log("train", T, stats)
        stats["T"] = T
        return stats

    def load(self, load_file):
        data = torch.load(load_file)
        self.online_actor.load_state_dict(data["online_actor"])
        self.online_Q1.load_state_dict(data["online_Q1"])
        self.online_Q2.load_state_dict(data["online_Q2"])
        self.target_actor.load_state_dict(data["target_actor"])
        self.target_Q1.load_state_dict(data["target_Q1"])
        self.target_Q2.load_state_dict(data["target_Q2"])
        self.actor_optimizer.load_state_dict(data["actor_optimizer"])
        self.Q_optimizer.load_state_dict(data["Q_optimizer"])

    def save(self, save_dir, file_name):
        data = {"online_actor": self.online_actor.state_dict(),
                "online_Q1": self.online_Q1.state_dict(),
                "online_Q2": self.online_Q2.state_dict(),
                "target_actor": self.target_actor.state_dict(),
                "target_Q1": self.target_Q1.state_dict(),
                "target_Q2": self.target_Q2.state_dict(),
                "actor_optimizer": self.actor_optimizer.state_dict(),
                "Q_optimizer": self.Q_optimizer.state_dict()}
        torch.save(data, os.path.join(save_dir, file_name))


import torch
import torch.nn as nn
import numpy as np

def predict(model, model_dir, results_dir, features_path, vid_list_file, epoch, actions_dict, device, sample_rate, args):

    # collect arguments
    verbose = args.verbose
    use_best_model = args.use_best_model

    # multi-GPU
    if args.multi_gpu and torch.cuda.device_count() > 1:
        model = nn.DataParallel(model)
    model.eval()

    with torch.no_grad():
        model.to(device)
        if use_best_model == 'source':
            model.load_state_dict(torch.load(model_dir + "/acc_best_source.model"))
        elif use_best_model == 'target':
            model.load_state_dict(torch.load(model_dir + "/acc_best_target.model"))
        else:
            model.load_state_dict(torch.load(model_dir + "/epoch-" + str(epoch) + ".model"))
        file_ptr = open(vid_list_file, 'r')
        list_of_vids = file_ptr.read().split('\n')[:-1]  # testing list
        file_ptr.close()
        for vid in list_of_vids:
            if verbose:
                print(vid)

            features = np.load(features_path + vid.split('.')[0] + '.npy')
            features = features[:, ::sample_rate]
            input_x = torch.tensor(features, dtype=torch.float)
            input_x.unsqueeze_(0)
            input_x = input_x.to(device)
            mask = torch.ones_like(input_x)
            predictions, _, _, _, _, _, _, _, _, _, _, _, _, _ = model(input_x, input_x, mask, mask, [0, 0], reverse=False)
            _, predicted = torch.max(predictions[:, -1, :, :].data, 1)
            predicted = predicted.squeeze()
            recognition = []
            for i in range(predicted.size(0)):
                recognition = np.concatenate((recognition,
                    [list(actions_dict.keys())[list(actions_dict.values()).index(predicted[i].item())]] * sample_rate))
            f_name = vid.split('/')[-1].split('.')[0]
            f_ptr = open(results_dir + "/" + f_name, "w")
            f_ptr.write("### Frame level recognition: ###\n")
            f_ptr.write(' '.join(recognition))
            f_ptr.close()


# type: ignore
import biodatasets
import numpy as np
from biotransformers import BioTransformers

data = biodatasets.load_dataset("swissProt")
X, y = data.to_npy_arrays(input_names=["sequence"])
X = X[0]

# Train on small sequence
length = np.array(list(map(len, X))) < 200
train_seq = X[length][:1000]
# print("Number of proteins {}".format(sum(length)))

def print_accuracy(acc_score):
    print("ACCURACY")
    print("Accuracy score : {}".format(acc_score))
    print("--------------")

bio_trans = BioTransformers("esm1_t6_43M_UR50S", device="cuda", multi_gpu=False)

acc_before = bio_trans.compute_accuracy(
    "/home/a.delfosse/data_swissprot.fasta", batch_size=16
)
print_accuracy(acc_before)

# bio_trans.train_masked(
#    train_seq,
#    toks_per_batch=500,
#    epochs=2,
#    accelerator="dp",
# )

# bio_trans.load_model("esm1_t6_43M_UR50S_finetuned.pt")

acc_after = bio_trans.compute_accuracy(train_seq, batch_size=16)
print_accuracy(acc_after)


'''
Twin Delayed DDPG (TD3), if no twin no delayed then it's DDPG.
using target Q instead of V net: 2 Q net, 2 target Q net, 1 policy net, 1 target policy net
original paper: https://arxiv.org/pdf/1802.09477.pdf
'''
import math
import random

import gym
import numpy as np

import torch
torch.multiprocessing.set_start_method('forkserver', force=True) # critical for make multiprocessing work
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
from torch.distributions import Normal

# from IPython.display import clear_output
import matplotlib.pyplot as plt
from matplotlib import animation
from IPython.display import display

import argparse
import time
import queue

import torch.multiprocessing as mp
from torch.multiprocessing import Process

from multiprocessing import Process, Manager
from multiprocessing.managers import BaseManager

from sim2real_policies.utils.envs import make_env
from sim2real_policies.utils.choose_env import choose_env
from sim2real_policies.utils.optimizers import SharedAdam, ShareParameters
from sim2real_policies.utils.buffers import ReplayBuffer
from sim2real_policies.utils.value_networks import *
from sim2real_policies.utils.policy_networks import DPG_PolicyNetwork
from sim2real_policies.utils.evaluate import evaluate
from sim2real_policies.utils.load_params import load_params

from mujoco_py import MujocoException

class TD3_Trainer():
    def __init__(self, replay_buffer, state_space, action_space, hidden_dim, q_lr, policy_lr,\
        action_range, policy_target_update_interval=1):
        self.replay_buffer = replay_buffer
        self.hidden_dim = hidden_dim

        self.q_net1 = QNetwork(state_space, action_space, hidden_dim)
        self.q_net2 = QNetwork(state_space, action_space, hidden_dim)
        self.target_q_net1 = QNetwork(state_space, action_space, hidden_dim)
        self.target_q_net2 = QNetwork(state_space, action_space, hidden_dim)
        self.policy_net = DPG_PolicyNetwork(state_space, action_space, hidden_dim, action_range)
        self.target_policy_net = DPG_PolicyNetwork(state_space, action_space, hidden_dim, action_range)
        print('Q Network (1,2): ', self.q_net1)
        print('Policy Network: ', self.policy_net)

        self.target_q_net1 = self.target_ini(self.q_net1, self.target_q_net1)
        self.target_q_net2 = self.target_ini(self.q_net2, self.target_q_net2)
        self.target_policy_net = self.target_ini(self.policy_net, self.target_policy_net)

        self.update_cnt = 0
        self.policy_target_update_interval = policy_target_update_interval

        self.q_optimizer1 = SharedAdam(self.q_net1.parameters(), lr=q_lr)
        self.q_optimizer2 = SharedAdam(self.q_net2.parameters(), lr=q_lr)
        self.policy_optimizer = SharedAdam(self.policy_net.parameters(), lr=policy_lr)

    def to_cuda(self):
        self.q_net1 = self.q_net1.cuda()
        self.q_net2 = self.q_net2.cuda()
        self.target_q_net1 = self.target_q_net1.cuda()
        self.target_q_net2 = self.target_q_net2.cuda()
        self.policy_net = self.policy_net.cuda()
        self.target_policy_net = self.target_policy_net.cuda()

    def target_ini(self, net, target_net):
        for target_param, param in zip(target_net.parameters(), net.parameters()):
            target_param.data.copy_(param.data)
        return target_net

    def target_soft_update(self, net, target_net, soft_tau):
    # Soft update the target net
        for target_param, param in zip(target_net.parameters(), net.parameters()):
            target_param.data.copy_(  # copy data value into target parameters
                target_param.data * (1.0 - soft_tau) + param.data * soft_tau
            )

        return target_net

    def update(self, batch_size, eval_noise_scale, reward_scale=10., gamma=0.9,soft_tau=1e-2):
        state, action, reward, next_state, done = self.replay_buffer.sample(batch_size)
        # print('sample:', state, action,  reward, 2done)

        state      = torch.FloatTensor(state).cuda()
        next_state = torch.FloatTensor(next_state).cuda()
        action     = torch.FloatTensor(action).cuda()
        reward     = torch.FloatTensor(reward).unsqueeze(1).cuda()  # reward is single value, unsqueeze() to add one dim to be [reward] at the sample dim;
        done       = torch.FloatTensor(np.float32(done)).unsqueeze(1).cuda()

        predicted_q_value1 = self.q_net1(state, action)
        predicted_q_value2 = self.q_net2(state, action)
        new_action = self.policy_net.evaluate(state, noise_scale=0.0)  # no noise, deterministic policy gradients
        new_next_action = self.target_policy_net.evaluate(next_state, noise_scale=eval_noise_scale) # clipped normal noise

        reward = reward_scale * (reward - reward.mean(dim=0)) / (reward.std(dim=0) + 1e-6) # normalize with batch mean and std; plus a small number to prevent numerical problem

    # Training Q Function
        target_q_min = torch.min(self.target_q_net1(next_state, new_next_action),self.target_q_net2(next_state, new_next_action))

        target_q_value = reward + (1 - done) * gamma * target_q_min # if done==1, only reward

        q_value_loss1 = ((predicted_q_value1 - target_q_value.detach())**2).mean()  # detach: no gradients for the variable
        q_value_loss2 = ((predicted_q_value2 - target_q_value.detach())**2).mean()
        self.q_optimizer1.zero_grad()
        q_value_loss1.backward()
        self.q_optimizer1.step()
        self.q_optimizer2.zero_grad()
        q_value_loss2.backward()
        self.q_optimizer2.step()

        if self.update_cnt%self.policy_target_update_interval==0:
            # Training Policy Function
            ''' implementation 1 '''
            # predicted_new_q_value = torch.min(self.q_net1(state, new_action),self.q_net2(state, new_action))
            ''' implementation 2 '''
            predicted_new_q_value = self.q_net1(state, new_action)

            policy_loss = - predicted_new_q_value.mean()
            self.policy_optimizer.zero_grad()
            policy_loss.backward()
            self.policy_optimizer.step()

        # Soft update the target nets
            self.target_q_net1=self.target_soft_update(self.q_net1, self.target_q_net1, soft_tau)
            self.target_q_net2=self.target_soft_update(self.q_net2, self.target_q_net2, soft_tau)
            self.target_policy_net=self.target_soft_update(self.policy_net, self.target_policy_net, soft_tau)

        self.update_cnt+=1

        return predicted_q_value1.mean()

    def save_model(self, path):
        torch.save(self.q_net1.state_dict(), path+'_q1')
        torch.save(self.q_net2.state_dict(), path+'_q2')
        torch.save(self.policy_net.state_dict(), path+'_policy')

    def load_model(self, path):
        self.q_net1.load_state_dict(torch.load(path+'_q1', map_location='cuda:0'))
        self.q_net2.load_state_dict(torch.load(path+'_q2', map_location='cuda:0'))
        self.policy_net.load_state_dict(torch.load(path+'_policy', map_location='cuda:0'))
        self.q_net1.eval()
        self.q_net2.eval()
        self.policy_net.eval()

    def share_memory(self):
        self.q_net1.share_memory()
        self.q_net2.share_memory()
        self.target_q_net1.share_memory()
        self.target_q_net2.share_memory()
        self.policy_net.share_memory()
        self.target_policy_net.share_memory()
        ShareParameters(self.q_optimizer1)
        ShareParameters(self.q_optimizer2)
        ShareParameters(self.policy_optimizer)

def worker(id, td3_trainer, env_name, environment_params, environment_wrappers, environment_wrapper_arguments, rewards_queue, eval_rewards_queue, success_queue,\
            eval_success_queue, eval_interval, replay_buffer, max_episodes, max_steps, batch_size, explore_steps, noise_decay, update_itr, explore_noise_scale, \
            eval_noise_scale, reward_scale, DETERMINISTIC, hidden_dim, model_path, seed=1):
    '''
    the function for sampling with multi-processing
    '''
    with torch.cuda.device(id % torch.cuda.device_count()):
        td3_trainer.to_cuda()
        print(td3_trainer, replay_buffer)
        env= make_env('robosuite.'+env_name, seed, id, environment_params, environment_wrappers,environment_wrapper_arguments)()
        action_dim = env.action_space.shape[0]
        frame_idx=0
        rewards=[]
        current_explore_noise_scale = explore_noise_scale
        # training loop
        for eps in range(max_episodes):
            episode_reward = 0
            state =  env.reset()
            current_explore_noise_scale = current_explore_noise_scale*noise_decay

            for step in range(max_steps):
                if frame_idx > explore_steps:
                    action = td3_trainer.policy_net.get_action(state, noise_scale=current_explore_noise_scale)
                else:
                    action = td3_trainer.policy_net.sample_action()

                try:
                    next_state, reward, done, info = env.step(action)
                    if environment_params["has_renderer"] and environment_params["render_visual_mesh"]:
                        env.render()
                except KeyboardInterrupt:
                    print('Finished')
                    td3_trainer.save_model(model_path)
                except MujocoException:
                    print('MujocoException')
                    break

                if info["unstable"]: # capture the case with cube flying away for pushing task
                    break

                replay_buffer.push(state, action, reward, next_state, done)

                state = next_state
                episode_reward += reward
                frame_idx += 1

                if replay_buffer.get_length() > batch_size:
                    for i in range(update_itr):
                        _=td3_trainer.update(batch_size, eval_noise_scale=eval_noise_scale, reward_scale=reward_scale)

                if done:
                    break
            print('Worker: ', id, '|Episode: ', eps, '| Episode Reward: ', episode_reward)
            rewards_queue.put(episode_reward)
            success_queue.put(info['success'])

            if eps % eval_interval == 0 and eps>0:
                # plot(rewards, id)
                td3_trainer.save_model(model_path)
                eval_r, eval_succ = evaluate(env, td3_trainer.policy_net)
                eval_rewards_queue.put(eval_r)
                eval_success_queue.put(eval_succ)

        td3_trainer.save_model(model_path)

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Train or test neural net motor controller.')
    parser.add_argument('--train', dest='train', action='store_true', default=False)
    parser.add_argument('--test', dest='test', action='store_true', default=False)

    args = parser.parse_args()

    # hyper-parameters for RL training
    env_name = ['SawyerReach', 'SawyerPush', 'SawyerSlide'][1]  # 36000 episodes for SawyerPush, 28000 episodes for SawyerReach and SawyerSlide
    max_episodes  = 9000
    max_steps   = 80   # it doesn't matter as long as it's larger than env.horizon
    num_workers = 4 # or: mp.cpu_count()
    eval_interval = 100  # evaluate training processs every N episodes
    # load other default parameters
    [action_range, batch_size, explore_steps, update_itr, explore_noise_scale, eval_noise_scale, reward_scale, \
        hidden_dim, noise_decay, policy_target_update_interval, q_lr, policy_lr, replay_buffer_size, DETERMINISTIC] = \
            load_params('td3', ['action_range', 'batch_size', 'explore_steps', 'update_itr', 'explore_noise_scale',\
             'eval_noise_scale', 'reward_scale', 'hidden_dim', 'noise_decay', \
                 'policy_target_update_interval', 'q_lr', 'policy_lr','replay_buffer_size', 'deterministic'] )

    # the replay buffer is a class, have to use torch manager to make it a proxy for sharing across processes
    BaseManager.register('ReplayBuffer', ReplayBuffer)
    manager = BaseManager()
    manager.start()
    replay_buffer = manager.ReplayBuffer(replay_buffer_size)  # share the replay buffer through manager

    env, environment_params, environment_wrappers, environment_wrapper_arguments = choose_env(env_name)
    prefix=env_name+str(len(environment_params["parameters_to_randomise"]))  # number of randomised parameters
    model_path = '../../../../data/td3/model/'+prefix+'_td3'

    action_space = env.action_space
    state_space = env.observation_space

    td3_trainer=TD3_Trainer(replay_buffer, state_space, action_space, hidden_dim, q_lr, policy_lr,\
         policy_target_update_interval=policy_target_update_interval, action_range=action_range )

    if args.train:
        # curriculum learning, pre-train with non-randomised
        # prefix=env_name+str(0)  # number of randomised parameters
        # ini_model_path = '../../../../data/td3/model/'+prefix+'_td3'
        # td3_trainer.load_model(ini_model_path)

        td3_trainer.load_model(model_path)
        td3_trainer.share_memory()

        rewards_queue=mp.Queue()  # used for get rewards from all processes and plot the curve
        eval_rewards_queue = mp.Queue()  # used for get offline evaluated rewards from all processes and plot the curve
        success_queue = mp.Queue()  # used for get success events from all processes
        eval_success_queue = mp.Queue()

        processes=[]
        rewards=[]
        success = []
        eval_rewards = []
        eval_success = []

        for i in range(num_workers):
            process = Process(target=worker, args=(i, td3_trainer, env_name, environment_params, environment_wrappers,environment_wrapper_arguments, \
            rewards_queue, eval_rewards_queue, success_queue, eval_success_queue, eval_interval, replay_buffer, max_episodes, max_steps, batch_size, explore_steps, noise_decay,\
            update_itr, explore_noise_scale, eval_noise_scale, reward_scale, DETERMINISTIC, hidden_dim, model_path))  # the args contain shared and not shared
            process.daemon=True  # all processes closed when the main stops
            processes.append(process)

        [p.start() for p in processes]
        while True:  # keep getting the episode reward from the queue
            # r = rewards_queue.get()
            # succ = success_queue.get()
            eval_r = eval_rewards_queue.get() # this queue has different sample frequence with above two queues, .get() at same time will break the while loop
            eval_succ = eval_success_queue.get()

            # success.append(succ)
            # rewards.append(r)
            eval_rewards.append(eval_r)
            eval_success.append(eval_succ)
            if len(eval_rewards)%20==0 and len(eval_rewards)>0:
                # plot(rewards)
                # np.save(prefix+'td3_rewards', rewards)
                # np.save(prefix+'td3_success', success)
                np.save('eval_rewards', eval_rewards)
                np.save('eval_success', eval_success)

        [p.join() for p in processes]  # finished at the same time

        td3_trainer.save_model(model_path)

    if args.test:
        td3_trainer.load_model(model_path)
        td3_trainer.to_cuda()
        env.renderer_on()
        for eps in range(10):
            state =  env.reset()
            env.render()
            episode_reward = 0

            for step in range(max_steps):
                action = td3_trainer.policy_net.get_action(state, noise_scale=0.0)
                next_state, reward, done, _ = env.step(action)
                env.render()

                episode_reward += reward
                state=next_state

                if done:
                    break

            print('Episode: ', eps, '| Episode Reward: ', episode_reward)


# This source code is part of the Hydride package and is distributed
# under the 3-Clause BSD License. Please see 'LICENSE.rst' for further
# information.

__name__ = "hydride"
__author__ = "Patrick Kunzmann"
__all__ = ["FragmentLibrary"]

from os.path import join, dirname, abspath
import warnings
import pickle
from biotite.structure.error import BadStructureError
from biotite.structure import BondType, displacement
import numpy as np

class FragmentLibrary:
    """
    A molecule fragment library for estimation of hydrogen positions.

    For each molecule added to the :class:`FragmentLibrary`,
    the molecule is split into fragments.
    Each fragment consists of

        - A central heavy atom,
        - bond order and position of its bonded heavy atoms and
        - and positions of bonded hydrogen atoms.

    The properties of the fragment (central atom element,
    central atom charge, order of connected bonds) are stored in
    a dictionary mapping these properties to heavy and hydrogen atom
    positions.

    If hydrogen atoms should be added to a target structure,
    the target structure is also split into fragments.
    Now the corresponding reference fragment in the library dictionary
    is accessed for each fragment.
    The corresponding atom coordinates of the reference fragment
    are superimposed [1]_ [2]_ onto the target fragment to obtain the
    hydrogen coordinates for the heavy atom.

    The constructor of this class creates an empty library.

    References
    ----------

    .. [1] W Kabsch,
       "A solution for the best rotation to relate two sets of vectors."
       Acta Cryst, 32, 922-923 (1976).

    .. [2] W Kabsch,
       "A discussion of the solution for the best rotation to relate
       two sets of vectors."
       Acta Cryst, 34, 827-828 (1978).
    """

    _std_library = None

    def __init__(self):
        self._frag_dict = {}

    @staticmethod
    def standard_library():
        """
        Get the standard :class:`FragmentLibrary`.
        The library contains fragments from all molecules in the
        *RCSB* *Chemical Component Dictionary*.

        Returns
        -------
        library : FragmentLibrary
            The standard library.
        """
        if FragmentLibrary._std_library is None:
            FragmentLibrary._std_library = FragmentLibrary()
            file_name = join(dirname(abspath(__file__)), "fragments.pickle")
            with open(file_name, "rb") as fragments_file:
                FragmentLibrary._std_library._frag_dict \
                    = pickle.load(fragments_file)
        return FragmentLibrary._std_library

    def add_molecule(self, molecule):
        """
        Add the fragments of a molecule to the library.

        Parameters
        ----------
        molecule : AtomArray
            A molecule, whose fragments should be added to the library.
            The structure must contain hydrogen atoms for each
            applicable heavy atom.
            The molecule must have an associated :class:`BondList`.
            The molecule must also include the *charge* annotation
            array, depicting the formal charge for each atom.
        """
        fragments = _fragment(molecule)
        for i, fragment in enumerate(fragments):
            if fragment is None:
                continue
            (
                central_element, central_charge, stereo, bond_types,
                center_coord, heavy_coord, hydrogen_coord
            ) = fragment
            # Translate the coordinates,
            # so the central heavy atom is at origin
            centered_heavy_coord = heavy_coord - center_coord
            centered_hydrogen_coord = hydrogen_coord - center_coord
            self._frag_dict[
                (central_element, central_charge, stereo, tuple(bond_types))
            ] = (
                # Information about the origin of the fragment
                # # for debugging purposes
                molecule.res_name[i], molecule.atom_name[i],
                # The interesting information
                centered_heavy_coord, centered_hydrogen_coord
            )
            if stereo != 0:
                # Also include the opposite enantiomer in the library
                # by reflecting the coordinates along an arbitrary axis
                stereo *= -1
                refl_centered_heavy_coord = centered_heavy_coord.copy()
                refl_centered_hydrogen_coord = centered_hydrogen_coord.copy()
                refl_centered_heavy_coord[..., 0] *= -1
                refl_centered_hydrogen_coord[..., 0] *= -1
                self._frag_dict[(
                    central_element, central_charge, stereo, tuple(bond_types)
                )] = (
                    molecule.res_name[i], molecule.atom_name[i],
                    refl_centered_heavy_coord, refl_centered_hydrogen_coord
                )

    def calculate_hydrogen_coord(self, atoms, mask=None, box=None):
        """
        Estimate the hydrogen coordinates for each atom in a given
        structure/molecule.

        Parameters
        ----------
        atoms : AtomArray, shape=(n,)
            The structure to get the hydrogen atom positions for.
            The structure must not contain any hydrogen atoms.
            The structure must have an associated :class:`BondList`.
            The structure must also include the *charge* annotation
            array, depicting the formal charge for each atom.
        mask : ndarray, shape=(n,), dtype=bool
            A boolean mask that is true for each heavy atom, where
            corresponsing hydrogen atom positions should be calculated.
            By default, hydrogen atoms are calculated for all applicable
            atoms.
        box : bool or array-like, shape=(3,3), dtype=float, optional
            If this parameter is set, periodic boundary conditions are
            taken into account (minimum-image convention), based on
            the box vectors given with this parameter.
            If `box` is set to true, the box vectors are taken from the
            ``box`` attribute of `atoms` instead.

        Returns
        -------
        hydrogen_coord : list of (ndarray, shape=(k,3), dtype=np.float32), length=n
            A list of hydrogen coordinates for each atom in the input
            `atoms`.
            *k* is the number of hydrogen atoms for this atom.
            Each atom, that is not included in the input `mask` or which
            has no fitting fragment in the library,
            has no corresponding hydrogen coordinates, i.e. *k* is 0.
        """
        if mask is None:
            mask = np.ones(atoms.array_length(), dtype=bool)

        # The target and reference heavy atom coordinates
        # for each fragment
        tar_frag_center_coord = np.zeros(
            (atoms.array_length(), 3), dtype=np.float32
        )
        tar_frag_heavy_coord = np.zeros(
            (atoms.array_length(), 3, 3), dtype=np.float32
        )
        ref_frag_heavy_coord = np.zeros(
            (atoms.array_length(), 3, 3), dtype=np.float32
        )
        # The amount of hydrogens varies for each fragment
        # -> padding with NaN
        # The maximum number of bond hydrogen atoms is 4
        ref_frag_hydrogen_coord = np.full(
            (atoms.array_length(), 4, 3), np.nan, dtype=np.float32
        )

        # Fill the coordinate arrays
        fragments = _fragment(atoms, mask)
        for i, fragment in enumerate(fragments):
            if fragment is None:
                # This atom is not in mask
                continue
            (
                central_element, central_charge, stereo, bond_types,
                center_coord, heavy_coord, _
            ) = fragment
            tar_frag_center_coord[i] = center_coord
            tar_frag_heavy_coord[i] = heavy_coord
            # The hydrogen_coord can be ignored:
            # In the target structure are no hydrogen atoms
            hit = self._frag_dict.get(
                (central_element, central_charge, stereo, tuple(bond_types))
            )
            if hit is None:
                warnings.warn(
                    f"Missing fragment for atom '{atoms.atom_name[i]}' "
                    f"at position {i}"
                )
            else:
                _, _, ref_heavy_coord, ref_hydrogen_coord = hit
                ref_frag_heavy_coord[i] = ref_heavy_coord
                ref_frag_hydrogen_coord[i, :len(ref_hydrogen_coord)] \
                    = ref_hydrogen_coord

        # Translate the target coordinates,
        # so the central heavy atom is at origin
        # This has already been done for the reference atoms
        # in the 'add_molecule()' method
        if box is not None:
            # Find shortest possible displacement vector for each heavy
            # atom according to minimum image convention
            if box is True:
                if atoms.box is None:
                    raise ValueError("Input structure has no associated box")
                box = atoms.box
            else:
                # Box vectors are given as array-like object
                box = np.asarray(box)
        tar_frag_heavy_coord = displacement(
            tar_frag_center_coord[:, np.newaxis, :], tar_frag_heavy_coord,
            box
        )

        # Get the rotation matrix required for superimposition of
        # the reference coord to the target coord
        matrices = _get_rotation_matrices(
            tar_frag_heavy_coord, ref_frag_heavy_coord
        )
        # Rotate the reference hydrogen atoms, so they fit the
        # target heavy atoms
        tar_frag_hydrogen_coord = _rotate(ref_frag_hydrogen_coord, matrices)
        # Translate hydrogen atoms to the position of the
        # non-centered central heavy target atom
        tar_frag_hydrogen_coord += tar_frag_center_coord[:, np.newaxis, :]

        # Turn into list and remove NaN paddings
        tar_frag_hydrogen_coord = [
            # If the x-coordinate is NaN it is expected that
            # y and z are also NaN
            coord[~np.isnan(coord[:, 0])] for coord in tar_frag_hydrogen_coord
        ]

        return tar_frag_hydrogen_coord

def _fragment(atoms, mask=None):
    """
    Create fragments for the input structure/molecule.

    Parameters
    ----------
    atoms : AtomArray, shape=(n,)
        The structure to be fragmented.
        The structure must have an associated :class:`BondList`.
        The structure must also include the *charge* annotation
        array, depicting the formal charge for each atom.
    mask : ndarray, shape=(n,), dtype=bool
        A boolean mask that is true for each heavy atom for which a
        fragment should be created.

    Returns
    -------
    fragments : list of tuple(str, int, int, ndarray, ndarray, ndarray), length=n
        The fragments.
        The tuple elements are

            #. the central atom element,
            #. the central atom charge,
            #. the enantiomer for stereocenter
               (``-1`` and ``1`` based on a custom nomenclature),
            #. the :class:`BondType` for each bonded heavy atom,
            #. the coordinates of the central atom,
            #. 3 coordinates of bonded heavy atoms (includes padding
               values, if there are not enough heavy atoms),
            #. the coordinates of bonded hydrogen atoms.

        ``None`` for each atom not included by the `mask`.
    """
    if mask is None:
        mask = np.ones(atoms.array_length(), dtype=bool)

    if atoms.bonds is None:
        raise BadStructureError(
            "The input structure must have an associated BondList"
        )

    fragments = [None] * atoms.array_length()

    all_bond_indices, all_bond_types = atoms.bonds.get_all_bonds()
    elements = atoms.element
    charges = atoms.charge
    coord = atoms.coord

    for i in range(atoms.array_length()):
        if not mask[i]:
            continue

        if elements[i] == "H":
            # Only create fragments for heavy atoms
            continue
        bond_indices = all_bond_indices[i]
        bond_types = all_bond_types[i]
        bond_indices = bond_indices[bond_indices != -1]
        bond_types = bond_types[bond_types != -1]

        heavy_mask = (elements[bond_indices] != "H")
        heavy_indices = bond_indices[heavy_mask]
        heavy_types = bond_types[heavy_mask]
        if (heavy_types == BondType.ANY).any():
            warnings.warn(
                f"Atom '{atoms.atom_name[i]}' in '{atoms.res_name[i]}' has an "
                f"undefined bond type and is ignored"
            )
            continue

        # Order the bonded atoms by their bond types
        # to remove atom order dependency in the matching step
        order = np.argsort(heavy_types)
        heavy_indices = heavy_indices[order]
        heavy_types = heavy_types[order]

        hydrogen_mask = ~heavy_mask
        hydrogen_coord = coord[bond_indices[hydrogen_mask]]

        # Special handling of nitrogen as central atom:
        # There are cases where the free electron pair can form
        # a partial double bond.
        # Although the bond order is formally 1 in this case,
        # it would enforce planar hydrogen positionioning
        # Therefore, a partial double bond is handled as bond type 7
        if elements[i] == "N":
            for j, remote_index in enumerate(heavy_indices):
                if heavy_types[j] != 1:
                    # This handling only applies to single bonds
                    continue
                rem_bond_indices = all_bond_indices[remote_index]
                rem_bond_indices = rem_bond_indices[rem_bond_indices != -1]
                rem_bond_types = all_bond_types[remote_index]
                rem_bond_types = rem_bond_types[rem_bond_types != -1]
                for rem_rem_index, bond_type in zip(
                    rem_bond_indices, rem_bond_types
                ):
                    # If the adjacent atom has a double bond
                    # the partial double bond condition is fulfilled
                    if bond_type == BondType.AROMATIC_DOUBLE or \
                       bond_type == BondType.DOUBLE:
                            heavy_types[j] = 7

        n_heavy_bonds = np.count_nonzero(heavy_mask)
        if n_heavy_bonds == 0:
            # The orientation is arbitrary
            # -> The fragment coord is the coord of the central atom
            # 4 times repeated
            heavy_coord = np.repeat(coord[np.newaxis, i, :], 3, axis=0)
            stereo = 0
        elif n_heavy_bonds == 1:
            # Include one atom further away
            # to get an unambiguous fragment
            remote_index = heavy_indices[0]
            rem_bond_indices = all_bond_indices[remote_index]
            rem_bond_indices = rem_bond_indices[rem_bond_indices != -1]
            rem_heavy_mask = (elements[rem_bond_indices] != "H")
            rem_heavy_indices = rem_bond_indices[rem_heavy_mask]
            # Use the coord of any heavy atom bonded to the remote
            # atom
            rem_rem_index = rem_heavy_indices[0]
            # Include the directly bonded atom two times, to give it a
            # greater weight in superimposition
            heavy_coord = coord[[remote_index, remote_index, rem_rem_index]]
            stereo = 0
        elif n_heavy_bonds == 2:
            heavy_coord = coord[[heavy_indices[0], heavy_indices[1], i]]
            stereo = 0
        elif n_heavy_bonds == 3:
            heavy_coord = coord[heavy_indices]
            center = coord[i]
            # Determine the enantiomer of this stereocenter
            # For performance reasons, the result does not follow the
            # R/S nomenclature, but a custom -1/1 based one, which also
            # unambiguously identifies the enantiomer
            n = np.cross(heavy_coord[0] - center, heavy_coord[1] - center)
            stereo = int(np.sign(np.dot(heavy_coord[2] - center, n)))
        elif n_heavy_bonds == 4:
            # The fragment is irrelevant, as there is no bonded hydrogen
            # -> The fragment coord is the coord of the central atom
            # 4 times repeated
            heavy_coord = np.repeat(coord[np.newaxis, i, :], 3, axis=0)
            stereo = 0
        else:
            warnings.warn(
                f"Atom '{atoms.atom_name[i]}' in "
                f"'{atoms.res_name[i]}' has more than 4 bonds to "
                f"heavy atoms ({n_heavy_bonds}) and is ignored"
            )
            heavy_coord = np.repeat(coord[np.newaxis, i, :], 3, axis=0)
            hydrogen_coord = np.zeros((0, 3), dtype=np.float32)
            stereo = 0
        central_coord = coord[i]
        fragments[i] = (
            elements[i], charges[i], stereo, heavy_types,
            central_coord, heavy_coord, hydrogen_coord
        )
    return fragments

def _get_rotation_matrices(fixed, mobile):
    """
    Get the rotation matrices to superimpose the given mobile
    coordinates into the given fixed coordinates, minimizing the RMSD.

    Uses the *Kabsch* algorithm.

    Parameters
    ----------
    fixed : ndarray, shape=(m,n,3), dtype=np.float32
        The fixed coordinates.
    mobile : ndarray, shape=(m,n,3), dtype=np.float32
        The mobile coordinates.

    Returns
    -------
    matrices : ndarray, shape=(m,3,3), dtype=np.float32
        The rotation matrices.
    """
    # Calculate cross-covariance matrices
    cov = np.sum(fixed[:,:,:,np.newaxis] * mobile[:,:,np.newaxis,:], axis=1)
    v, s, w = np.linalg.svd(cov)
    # Remove possibility of reflected atom coordinates
    reflected_mask = (np.linalg.det(v) * np.linalg.det(w) < 0)
    v[reflected_mask, :, -1] *= -1
    matrices = np.matmul(v, w)
    return matrices

def _rotate(coord, matrices):
    """
    Apply a rotation on given coordinates.

    Parameters
    ----------
    coord : ndarray, shape=(m,n,3), dtype=np.float32
        The coordinates.
    matrices : ndarray, shape=(m,3,3), dtype=np.float32
        The rotation matrices.
    """
    return np.transpose(
        np.matmul(matrices, np.transpose(coord, axes=(0, 2, 1))),
        axes=(0, 2, 1)
    )


#!/usr/bin/env python3
# test.py
import time
import sys
# import mss
sys.path.append('../')
from mss import visreader as mvis
from mss import mssmain as msm
from mss import align
import pandas as pd
import numpy as np
from tqdm import tqdm
import peakutils
import scipy
from scipy.integrate import simps
import itertools
path = '../example_data/ex_1.mzML'
scans = msm.get_scans(path, ms_all=False, ms_lv=1)
#noise removal
msm.noise_removal(scans, 2000)
scans = scans[:200]

@profile
def ms_chromatogram_list(mzml_scans, input_mz, error):
    '''
    Generate a peak list for specific input_mz over
    whole rt period from the mzml file
    ***Most useful function!
    '''
    intensity = []
    retention_time = [i.scan_time[0] for i in mzml_scans]
    for scan in mzml_scans:
        _, target_index = msm.mz_locator(scan.mz, input_mz, error)
        if len(target_index) == 0:
            intensity.append(0)
        else:
            # intensity.append(scan.i[target_index])
            # CR -> if all_than_close=True
            # Change from sum to max
            intensity.append(max(scan.i[target_index]))

    return retention_time, intensity

ms_chromatogram_list(scans,299.1765,500)