from EXOSIMS.Prototypes.SurveySimulation import SurveySimulation
import astropy.units as u
import numpy as np
import astropy.constants as const


class linearJScheduler_sotoSS(SurveySimulation):
    """linearJScheduler

    This class implements the linear cost function scheduler described
    in Savransky et al. (2010).

    Args:
        coeffs (iterable 3x1):
            Cost function coefficients: slew distance, completeness, target list
            coverage

        **specs:
            user specified values

    """

    def __init__(self, coeffs=[1, 1, 2, 1], revisit_wait=91.25, **specs):

        SurveySimulation.__init__(self, **specs)

        # verify that coefficients input is iterable 4x1
        if not (isinstance(coeffs, (list, tuple, np.ndarray))) or (len(coeffs) != 4):
            raise TypeError("coeffs must be a 4 element iterable")

        # Add to outspec
        self._outspec["coeffs"] = coeffs
        self._outspec["revisit_wait"] = revisit_wait

        # normalize coefficients
        coeffs = np.array(coeffs)
        coeffs = coeffs / np.linalg.norm(coeffs)

        self.coeffs = coeffs

        self.revisit_wait = revisit_wait * u.d
        self.no_dets = np.ones(self.TargetList.nStars, dtype=bool)

    def choose_next_target(self, old_sInd, sInds, slewTimes, intTimes):
        """Choose next target based on truncated depth first search
        of linear cost function.

        Args:
            old_sInd (int):
                Index of the previous target star
            sInds (int array):
                Indices of available targets
            slewTimes (astropy quantity array):
                slew times to all stars (must be indexed by sInds)
            intTimes (astropy Quantity array):
                Integration times for detection in units of day

        Returns:
            tuple:
                sInd (int):
                    Index of next target star
                waitTime (astropy Quantity):
                    the amount of time to wait (this method returns None)

        """

        Comp = self.Completeness
        TL = self.TargetList
        TK = self.TimeKeeping
        OS = self.OpticalSystem
        Obs = self.Observatory
        allModes = OS.observingModes

        # cast sInds to array
        sInds = np.array(sInds, ndmin=1, copy=False)

        # current star has to be in the adjmat
        if (old_sInd is not None) and (old_sInd not in sInds):
            sInds = np.append(sInds, old_sInd)

        # calculate dt since previous observation
        dt = TK.currentTimeNorm.copy() + slewTimes[sInds] - self.lastObsTimes[sInds]
        # get dynamic completeness values
        comps = Comp.completeness_update(TL, sInds, self.starVisits[sInds], dt)

        # if first target, or if only 1 available target,
        # choose highest available completeness
        nStars = len(sInds)
        if (old_sInd is None) or (nStars == 1):
            sInd = np.random.choice(sInds[comps == max(comps)])
            return sInd, slewTimes[sInd]

        # define adjacency matrix
        A = np.zeros((nStars, nStars))

        # only consider slew distance when there's an occulter
        if OS.haveOcculter:
            r_ts = TL.starprop(sInds, TK.currentTimeAbs.copy())
            u_ts = (
                r_ts.to("AU").value.T / np.linalg.norm(r_ts.to("AU").value, axis=1)
            ).T
            angdists = np.arccos(np.clip(np.dot(u_ts, u_ts.T), -1, 1))
            A[np.ones((nStars), dtype=bool)] = angdists
            A = self.coeffs[0] * (A) / np.pi

        # add factor due to completeness
        A = A + self.coeffs[1] * (1 - comps)

        # add factor due to unvisited ramp
        f_uv = np.zeros(nStars)
        unvisited = self.starVisits[sInds] == 0
        f_uv[unvisited] = float(TK.currentTimeNorm.copy() / TK.missionLife.copy()) ** 2
        A = A - self.coeffs[2] * f_uv

        # add factor due to revisited ramp
        f2_uv = 1 - (np.in1d(sInds, self.starRevisit[:, 0]))
        A = A + self.coeffs[3] * f2_uv

        # kill diagonal
        A = A + np.diag(np.ones(nStars) * np.Inf)

        # take two traversal steps
        step1 = np.tile(A[sInds == old_sInd, :], (nStars, 1)).flatten("F")
        step2 = A[np.array(np.ones((nStars, nStars)), dtype=bool)]
        tmp = np.argmin(step1 + step2)
        sInd = sInds[int(np.floor(tmp / float(nStars)))]

        waitTime = slewTimes[sInd]
        # Check if exoplanetObsTime would be exceeded
        mode = list(filter(lambda mode: mode["detectionMode"], allModes))[0]
        (
            maxIntTimeOBendTime,
            maxIntTimeExoplanetObsTime,
            maxIntTimeMissionLife,
        ) = TK.get_ObsDetectionMaxIntTime(Obs, mode)
        maxIntTime = min(
            maxIntTimeOBendTime, maxIntTimeExoplanetObsTime, maxIntTimeMissionLife
        )  # Maximum intTime allowed
        intTimes2 = self.calc_targ_intTime(sInd, TK.currentTimeAbs.copy(), mode)
        if (
            intTimes2 > maxIntTime
        ):  # check if max allowed integration time would be exceeded
            self.vprint("max allowed integration time would be exceeded")
            sInd = None
            waitTime = 1.0 * u.d

        return sInd, waitTime

    def revisitFilter(self, sInds, tmpCurrentTimeNorm):
        """Helper method for Overloading Revisit Filtering

        Args:
            sInds - indices of stars still in observation list
            tmpCurrentTimeNorm (MJD) - the simulation time after overhead was added in
            MJD form

        Returns:
            ~numpy.ndarray(int):
                indices of stars still in observation list
        """
        tovisit = np.zeros(
            self.TargetList.nStars, dtype=bool
        )  # tovisit is a boolean array containing the
        if len(sInds) > 0:  # so long as there is at least 1 star left in sInds
            tovisit[sInds] = (self.starVisits[sInds] == min(self.starVisits[sInds])) & (
                self.starVisits[sInds] < self.nVisitsMax
            )  # Checks that no star has exceeded the number of revisits
            if (
                self.starRevisit.size != 0
            ):  # There is at least one revisit planned in starRevisit
                dt_rev = (
                    self.starRevisit[:, 1] * u.day - tmpCurrentTimeNorm
                )  # absolute temporal spacing between revisit and now.

                # return indices of all revisits within a threshold dt_max of revisit
                # day and indices of all revisits with no detections past the revisit
                # time
                ind_rev2 = [
                    int(x)
                    for x in self.starRevisit[dt_rev < 0 * u.d, 0]
                    if (x in sInds)
                ]
                tovisit[ind_rev2] = self.starVisits[ind_rev2] < self.nVisitsMax
            sInds = np.where(tovisit)[0]

        return sInds

    def scheduleRevisit(self, sInd, smin, det, pInds):
        """A Helper Method for scheduling revisits after observation detection
        Args:
            sInd - sInd of the star just detected
            smin - minimum separation of the planet to star of planet just detected
            det -
            pInds - Indices of planets around target star
        Return:
            updates self.starRevisit attribute
        """
        TK = self.TimeKeeping
        TL = self.TargetList
        SU = self.SimulatedUniverse

        # in both cases (detection or false alarm), schedule a revisit
        # based on minimum separation
        Ms = TL.MsTrue[sInd]
        if (
            smin is not None and smin is not np.nan
        ):  # smin is None if no planet was detected
            sp = smin
            if np.any(det):
                pInd_smin = pInds[det][np.argmin(SU.s[pInds[det]])]
                Mp = SU.Mp[pInd_smin]
            else:
                Mp = SU.Mp.mean()
            mu = const.G * (Mp + Ms)
            T = 2.0 * np.pi * np.sqrt(sp**3 / mu)
            t_rev = TK.currentTimeNorm.copy() + T / 2.0
        # otherwise, revisit based on average of population semi-major axis and mass
        else:
            sp = SU.s.mean()
            Mp = SU.Mp.mean()
            mu = const.G * (Mp + Ms)
            T = 2.0 * np.pi * np.sqrt(sp**3 / mu)
            t_rev = TK.currentTimeNorm.copy() + 0.75 * T

        # if no detections then schedule revisit based off of revisit_weight
        if not np.any(det):
            t_rev = TK.currentTimeNorm.copy() + self.revisit_wait
            self.no_dets[sInd] = True
        else:
            self.no_dets[sInd] = False

        t_rev = TK.currentTimeNorm.copy() + self.revisit_wait
        # finally, populate the revisit list (NOTE: sInd becomes a float)
        revisit = np.array([sInd, t_rev.to("day").value])
        if self.starRevisit.size == 0:  # If starRevisit has nothing in it
            self.starRevisit = np.array([revisit])  # initialize sterRevisit
        else:
            revInd = np.where(self.starRevisit[:, 0] == sInd)[
                0
            ]  # indices of the first column of the starRevisit list containing sInd
            if revInd.size == 0:
                self.starRevisit = np.vstack((self.starRevisit, revisit))
            else:
                self.starRevisit[revInd, 1] = revisit[1]  # over
