#!/usr/bin/env python

"""Basis for developing a scheduler for tracking multiple space objects."""
import logging
from abc import abstractmethod

import numpy as np
import pyorb

from ..controller_v1 import Tracker
from .scheduler import Scheduler
from ..passes import equidistant_sampling

logger = logging.getLogger(__name__)


class Tracking(Scheduler):
    """#TODO: Docstring"""

    def __init__(
        self,
        radar,
        space_objects,
        end_time,
        epoch,
        start_time=0.0,
        controller=Tracker,
        controller_args=dict(return_copy=True),
        max_dpos=1e3,
        max_samp=30.0,
        use_pass_states=True,
        calculate_max_snr=False,
        collect_passes=False,
    ):
        super().__init__(radar)
        self.epoch = epoch
        self.controller = controller
        self.controller_args = controller_args
        self.space_objects = space_objects
        self.start_time = start_time
        self.end_time = end_time
        self.collect_passes = collect_passes

        self.max_dpos = max_dpos
        self.max_samp = max_samp
        self.use_pass_states = use_pass_states
        self.calculate_max_snr = calculate_max_snr

        self.reset_space_objects(space_objects)

    def reset_space_objects(self, space_objects):
        self.space_objects = space_objects
        self.states = [None] * len(space_objects)
        self.states_t = [None] * len(space_objects)
        self.passes = [None] * len(space_objects)

        if self.use_pass_states:
            self.measurements = [None] * len(space_objects)
        else:
            self.measurements = None

    def add_space_object(self, space_object):
        self.space_objects.append(space_object)
        self.states.append(None)
        self.states_t.append(None)
        self.passes.append(None)

        if self.use_pass_states:
            self.measurements.append(None)

    def update(self):
        self.reset_space_objects(self.space_objects)
        for ind in range(len(self.space_objects)):
            self.get_passes(ind)

    def get_passes(self, ind):
        """#TODO: Docstring"""
        dt = (self.space_objects[ind].epoch - self.epoch).to_value("sec")

        if isinstance(self.space_objects[ind].state, pyorb.Orbit):
            t = equidistant_sampling(
                orbit=self.space_objects[ind].state,
                start_t=self.start_time,
                end_t=self.end_time,
                max_dpos=self.max_dpos,
            )
        else:
            t = np.arange(self.start_time, self.end_time, self.max_samp)

        logger.info(f"Tracking:get_passes(ind={ind}):propagating {len(t)} steps")

        states = self.space_objects[ind].get_state(t - dt)
        if self.use_pass_states:
            self.states[ind] = states
            self.states_t[ind] = t

        logger.info(f"Tracking:get_passes(ind={ind}):propagating complete")

        self.passes[ind] = self.radar.find_passes(t, states, cache_data=True)

        # we may need SNR to plan observations
        for txi in range(len(self.radar.tx)):
            for rxi in range(len(self.radar.rx)):
                logger.info(
                    f"Tracking:get_passes(ind={ind}):tx{txi}-rx{rxi} {len(self.passes[ind][txi][rxi])} passes"
                )
                if not self.calculate_max_snr:
                    continue
                for ps in self.passes[ind][txi][rxi]:
                    ps.calculate_max_snr(
                        rx=self.radar.rx[rxi],
                        tx=self.radar.tx[txi],
                        diameter=self.space_objects[ind].d,
                    )

        return self.passes[ind], states, t

    @abstractmethod
    def set_measurements(self, *args, **kwargs):
        """#TODO: Docstring"""
        pass

    def get_controllers(self):
        """#TODO: Docstring"""
        logger.debug(f"Tracking:get_controllers")

        ctrls = []
        for ind in range(len(self.space_objects)):
            if self.collect_passes:
                if self.use_pass_states:
                    minds = self.measurements[ind]
                    states = self.states[ind][:3, minds]
                    t = self.states_t[ind][minds]
                else:
                    states = self.states[ind][:3, :]
                    t = self.states_t[ind][:]

                ctrl = self.controller(
                    radar=self.radar, t=t, t0=0.0, ecefs=states[:3, :], **self.controller_args
                )
                ctrl.meta["target"] = f"Object {self.space_objects[ind].object_id}"
                ctrls.append(ctrl)

            else:
                for pass_ind in range(len(self.states_t[ind])):
                    if self.use_pass_states:
                        minds = self.measurements[ind][pass_ind]
                        states = self.states[ind][pass_ind][:3, minds]
                        t = self.states_t[ind][pass_ind][minds]
                    else:
                        states = self.states[ind][pass_ind][:3, :]
                        t = self.states_t[ind][pass_ind][:]

                    ctrl = self.controller(
                        radar=self.radar, t=t, t0=0.0, ecefs=states[:3, :], **self.controller_args
                    )
                    ctrl.meta["target"] = f"Object {self.space_objects[ind].object_id}"
                    ctrl.meta["pass"] = pass_ind
                    ctrls.append(ctrl)

        return ctrls


class PriorityTracking(Tracking):
    """#TODO: Docstring"""

    def __init__(self, radar, space_objects, end_time, epoch, timeslice, allocation, **kwargs):
        self.priority = kwargs.pop("priority", None)
        super().__init__(radar, space_objects, end_time, epoch, **kwargs)

        self.timeslice = timeslice
        self.allocation = allocation

        if self.priority is None:
            self.priority = np.ones((len(self.space_objects),), dtype=np.float64)

    def add_space_object(self, space_object, priority):
        super().add_space_object(space_object)
        self.priority.append(priority)

    def update(self, space_objects=None, priority=None):
        if space_objects is not None:
            self.space_objects = space_objects
        if priority is not None:
            self.priority = priority
        assert len(self.priority) == len(self.space_objects), "Need a priority for each object"

        super().update()

    def set_measurements(self):
        """#TODO: Docstring"""
        if not self.collect_passes:
            raise NotImplementedError()

        total_measurements = int(self.allocation / self.timeslice)
        if not isinstance(self.priority, np.ndarray):
            pri = np.array(self.priority)
            pri = pri / pri.sum()
        else:
            pri = self.priority / self.priority.sum()

        # per object
        measurements = np.floor(pri * total_measurements).astype(np.int64)

        for ind in range(len(self.space_objects)):
            if self.use_pass_states:
                all_inds = []
                for txi in range(len(self.passes[ind])):
                    for rxi in range(len(self.passes[ind][txi])):
                        all_inds += [ps.inds for ps in self.passes[ind][txi][rxi]]
                if len(all_inds) == 0:
                    self.measurements[ind] = np.empty((0,), dtype=np.int64)
                    continue
                all_inds = np.concatenate(all_inds, axis=0)
                all_inds = np.unique(all_inds)

                # uniform distribution
                dt = int(len(all_inds) / measurements[ind])
                if dt == 0:
                    dt = 1
                all_inds = all_inds[::dt]
                self.measurements[ind] = all_inds
            else:
                raise NotImplementedError()
