from typing import Any, Dict, List, Optional, Union, TYPE_CHECKING

from theng.data.a_data_object import ADataObject
from theng.data.pathfinder import Trigger, Target
from theng.data.pathfinder.tags import Tag, TagCollection
from theng.data.point import Point
from theng.data.time_series import TimeSeries
from theng.data.vector import Vector
from theng.args import get_arguments


class Occupant(ADataObject):
    """Represents a single Occupant in a Pathfinder Simulation.

    Args:
        occ_id (str): Unique ID for this Occupant
        general_data (Dict): Dictionary of raw data from the _occupants.json file.
        param_data (Optional[Dict]): Optional Dictionary of raw data from the _occupant_params.json file
        detailed_data (Optional[Dict]): Optional Dictionary of raw data from the detailed occupant output files.
        sd_accumulated_data (Optional[Dict]): Optional Dictionary of raw data from the _sd_accumulated.json file.
        sd_transient_data (Optional[Dict]): Optional Dictionary of raw data from the _sd_transient.json file.
        parent (PathfinderModel): The PathfinderModel that this Occupant is a part of.
        
    Attributes:
        initial_params (OccParams): Initial parameters assigned to the Occupant in the Simulator. Contains the data of the _occupant_params files.
        history (Optional[TimeSeries[OccFrame]]): Optional TimeSeries detailed output data for this occupant. 
            None if the _occupant_detailed files are not present, or data is not present for this occupant.
        social_distancing (Optional[SocialDistanceData]): Optional SocialDistancingData for this occupant. 
            None if the _occupant_params files are not present, or data is not present for this occupant.
        
        id (int): A unique integer identifier for the occupant assigned by the simulator.
        name (str): The name given to the occupant in the user interface.
        exit_time (float): The time at which the Occupant went through and Exit Door.
            Default value is infinity. 
            This value represents that the Occupant has not exited the model.
        active_time (float): The amount of time the occupant was actively seeking a location in the model, such as a waypoint or exit.
        congestion_time_total (float):Total time the occupant spent moving at less than the congestion Velocity as specified in the simulation parameters (see `Section 14.1.2 of the Pathfinder User Manual <https://support.thunderheadeng.com/docs/pathfinder/latest/user-manual/#_output_parameters>`_).
        congestion_time_max_continuous (float): The maximum continuous amount of time the occupant spent at less than the congestion Velocity.
        congestion_time_level (float): Time spent by the occupant moving at less than congestion Velocity on level surfaces.
        congestion_time_stair (float): Time spent by the occupant moving at less than congestion Velocity on stairs.
        congestion_time_ramp (float): Time spent by the occupant moving at less than congestion Velocity on ramps.
        start_time (float): The time at which the occupant was generated. 
            For occupants pre-seeded in the model, this is always 0. 
            For occupants generated by an occupant source, this could be any time.
            The default value is infinity, representing that the Occupant did not report a start time.
        finish_time (float): The time at which the occupant left the simulation for any reason.
            The default value is infinity, representing that the occupant either did not exit the simulation, or did not report a finish time.
        distance (float): The total distance the occupant travelled during the simulation.
        tags (TagCollection): Collection of :class:`~theng.data.pathfinder.Tag` objects that were applied to this Occupant throughout the simulation.
        triggers_used (int): The total number of triggers used by the occupant during the simulation.
        triggers_usage_time (float): The amount of time spent by the occupant using triggers.
        targets_used (int): The total number of unique reserved Occupant Targets where the occupant waited during the simulation.
        target_usage_time (float): The total amount of time the occupant waited at a reserved Occupant Target.

        
        _general_data (Dict): Raw dict of data from the _occupants data files.
        _param_data (Dict): Raw dict of data from the _occupants_params data files.
        _detailed_data (Dict): Raw dict of data from the _occupants_detailed data files.
        _sd_accum_data (Dict): Raw dict of data from the _sd_accumulated data files.
        _sd_transient_data (Dict): Raw dict of data from the _sd_transient data files.
    """
    
    def __init__(
            self, 
            *, 
            occ_id: str, 
            general_data: Dict, 
            param_data: Optional[Dict], 
            detailed_data: Optional[Dict], 
            sd_accumulated_data: Optional[Dict], 
            sd_transient_data: Optional[Dict], 
            parent: Any = None
        ) -> None:
        super().__init__(parent=parent)
        self.id: int = int(occ_id)
        self._general_data: Dict = general_data
        self._param_data: Optional[Dict] = param_data
        self._detailed_data: Optional[Dict] = detailed_data
        self._sd_accum_data: Optional[Dict] = sd_accumulated_data
        self._sd_transient_data: Optional[Dict] = sd_transient_data
        
        # Loaded from General Occupant data
        self.name: str = ""
        self.exit_time: float = float('inf')
        self.active_time: float = 0.0
        self.congestion_time_total: float = 0.0
        self.congestion_time_max_continuous: float = 0.0
        self.congestion_time_level: float = 0.0
        self.congestion_time_stair: float = 0.0
        self.congestion_time_ramp: float = 0.0
        self.start_time: float = float('inf')
        self.finish_time: float = float('inf')
        self.distance: float = 0.0
        self.triggers_used: int = 0
        self.triggers_usage_time: float = 0.0
        self.targets_used: int = 0
        self.target_usage_time: float = 0.0
        self._tags_raw: Optional[Dict] = None
        self._tags: Dict[str, Tag] = {}
        self.tags: TagCollection = TagCollection([])
        
        # Loaded from Param data
        self.initial_params: OccParams = OccParams(param_data=None, parent=self)
        
        #Loaded from Detailed data
        self.history: Optional[TimeSeries[OccFrame]] = None
        
        #loaded from Transient and Accumulated Social Distancing Data
        self.social_distancing: Optional[SocialDistanceData] = None
        
        self._log.debug(f'Loading Occupant {self.id}')
        
        self._load_general_data()
        if self._param_data != None:
            self._load_param_data()
        if self._detailed_data != None:
            self._load_detailed_data()
        if self._sd_accum_data != None or self._sd_transient_data != None:
            self._load_sd_data()
        
    def _load_general_data(self) -> None:
        """Loads the data from the _occupants.json file"""
        self._log.debug(f'Loading General Occupant data for {self.id}')
        
        self._load_safe(source_data=self._general_data, data_key='name', target='name')
        self._load_safe(source_data=self._general_data, data_key='exitTime', target='exit_time')
        self._load_safe(source_data=self._general_data, data_key='activeTime', target='active_time')
        self._load_safe(source_data=self._general_data, data_key='congestionTimeTotal', target='congestion_time_total')
        self._load_safe(source_data=self._general_data, data_key='congestionTimeMaxContinuous', target='congestion_time_max_continuous')
        self._load_safe(source_data=self._general_data, data_key='congestionTimeLevel', target='congestion_time_level')
        self._load_safe(source_data=self._general_data, data_key='congestionTimeStair', target='congestion_time_stair')
        self._load_safe(source_data=self._general_data, data_key='congestionTimeRamp', target='congestion_time_ramp')
        self._load_safe(source_data=self._general_data, data_key='startTime', target='start_time')
        self._load_safe(source_data=self._general_data, data_key='finishTime', target='finish_time')
        self._load_safe(source_data=self._general_data, data_key='distance', target='distance')
        self._load_safe(source_data=self._general_data, data_key='numTriggersUsed', target='triggers_used')
        self._load_safe(source_data=self._general_data, data_key='triggersTime', target='triggers_usage_time')
        self._load_safe(source_data=self._general_data, data_key='numTargetsUsed', target='targets_used')
        self._load_safe(source_data=self._general_data, data_key='targetTime', target='target_usage_time')
        self._load_safe(source_data=self._general_data, data_key='tags', target='_tags_raw')
        
        if (self._tags_raw != None):
            self._load_tags_data()
            
    def _load_tags_data(self) -> None:
        """Format the loaded Tag data for this Occupant in to a TagCollection"""
        try:
            for key, data in self._tags_raw.items():
                self._tags[key] = Tag(name=key, data=data)
            self.tags = TagCollection(self._tags)
        except:
            self.tags = TagCollection([])
            if not get_arguments().quiet:
                self._log.warning(f'There was an error parsing tags for Occupant "{self.name}"')
        
    def _load_param_data(self) -> None:
        """Load the _occupant_params.json data"""
        self.initial_params = OccParams(param_data=self._param_data, parent=self)
        
    def _load_detailed_data(self) -> None:
        """load the detailed occupant output data"""        
        try:
            time: List[float] = []
            frames: List[OccFrame] = []
            for step, frame in self._detailed_data.items():
                time.append(float(step))
                frames.append(OccFrame(detailed_data=frame, parent=self))
            self.history = TimeSeries(time=time, values=frames)
        except:
            self.history = None
     
    def _load_sd_data(self) -> None:
        """Load social distancing data for this Occupant"""
        self.social_distancing = SocialDistanceData(sd_accum=self._sd_accum_data, sd_transient=self._sd_transient_data, parent=self)
        
    def update_occupant_references(self) -> None:
        """Update references to other Occupants in this Occupant object"""
        if self.social_distancing != None and type(self.social_distancing) == SocialDistanceData:
            self.social_distancing._update_occupant_references()
        
class OccParams(ADataObject):
    """Initial simulation parameters for an Occupant.

    Args:
        param_data (Optional[Dict]): Dict of raw param data for the Occupant.
        parent (Occupant): The Occupant associated with this data
        
    Attributes:
        parent (Occupant): The Occupant associated with this data
        profile (str): The name of the Profile assigned to this Occupant. 
        behavior (str): The name of the Behavior assigned to this Occupant.
        max_velocity (float): The maximum velocity allowed for this Occupant.
        shape (str): Name of the Occupant's assigned shape.
        room (str): The name of the room where the Occupant was created.
        position (Dict): Dict of X, Y, Z data of the Occupant's initial position.
        profile_props(Dict): Dictionary of profile properties. This is a temporary dictionary and will be changed in future versions
        
        _param_data (Optional[Dict]): Dict of raw param data for the Occupant.
    """
    
    def __init__(self, *, param_data: Optional[Dict], parent: Occupant):
        super().__init__()
        self._param_data: Optional[Dict] = param_data
        self.parent: Occupant = parent
        
        # Loaded from Occupant param data
        self.profile: str = ""
        self.behavior: str = ""
        self.max_velocity: float = float('inf')
        self.shape: str = {}
        self.room: str = ""
        self.position: Dict = {}
        self.profile_props: Optional[Dict] = None
        
        if self._param_data != None:
            self._load()
        
    def _load(self):
        """Loads the data from the _occupant_params.json file"""
        self._log.debug(f'Loading Occupant Parameter data for {self.parent.id}')
        
        self._load_safe(source_data=self._param_data, data_key="profile", target="profile")
        self._load_safe(source_data=self._param_data, data_key="behavior", target="behavior")
        self._load_safe(source_data=self._param_data, data_key="maxVelocity", target="max_velocity")
        self._load_safe(source_data=self._param_data, data_key="shape", target="shape")
        self._load_safe(source_data=self._param_data, data_key="room", target="room")
        self._load_safe(source_data=self._param_data, data_key="position", target="position")
        self._load_safe(source_data=self._param_data, data_key="profileProps", target="profile_props")
        
class OccFrame(ADataObject):
    """Represents the detailed data of an Occupant at a single time step.

    Args:
        detailed_data (Dict): The raw parsed detailed data for this frame of data.
        parent (Occupant): The :class:`~theng.data.pathfinder.Occupant` to which this frame of data applies.
        
    Attributes:
        parent (Occupant): The :class:`~theng.data.pathfinder.Occupant` to which this frame of data applies.
        active (bool): Whether the occupant is actively seeking an exit (True if they are seeking an exit and False if not).
        position (Point): The 3D location of the Occupant.
        velocity (Vector): The velocity of the Occupant.
        distance_traveled (float): The total distance the Occupant has travelled until this time in the simulation.
        location (str): The Occupant's current room.
        terrain (str): The terrain type of the Occupant's current room.
        trigger (Optioanl[Union[Trigger, str]]): The Trigger in use by the Occupant at the current time. None if no Trigger is in use.
        target (Optional[Union[Target, str]]): The Target in use by the Occupant at the current time. None if no Target is in use.
        co2_volume_fraction (Optional[float]): The current CO2 volume fraction exposure, if simulation is run with the FDS integration.
        co_volume_fraction (Optional[float]): The current CO volume fraction exposure, if simulation is run with the FDS integration.
        o2_volume_fraction (Optional[float]): The current O2 volume fraction exposure, if simulation is run with the FDS integration.
        temperature (Optional[float]): The current temperature exposure, if simulation is run with the FDS integration.
        soot_visibility (Optional[float]): The current visibility exposure, if simulation is run with the FDS integration.
        fed_total (Optional[float]): The current FED exposure, if simulation is run with the FDS integration.
    """
    
    def __init__(self, *, detailed_data: Dict, parent: Occupant) -> None:
        super().__init__()
        self._detailed_data: Dict = detailed_data
        self.parent: Occupant = parent
        
        self.active: bool = False
        self.position:  Point = Point()
        self.velocity: Vector = Vector()
        self.distance_traveled: float = float('inf')
        self.location: str = ""
        self.terrain: str = ""
        self.trigger: Optional[Union[Trigger, str]] = None #Trigger is loaded as a string from the file, then searched for in the model's TriggerCollection
        self.target: Optional[Union[Target, str]] = None #Target is loaded as a string from the file, then searched for in the model's TargetCollection
        self._tags_raw: List[str] = []
        self._tags: List[Tag] = []
        self.tags: TagCollection = TagCollection([])
        
        #FDS Integration data
        self.co2_volume_fraction: Optional[float] = None
        self.co_volume_fraction: Optional[float] = None
        self.o2_volume_fraction: Optional[float] = None
        self.temperature: Optional[float] = None
        self.soot_visibility: Optional[float] = None
        self.fed_total: Optional[float] = None
        
        self._load()
        
    def _load(self) -> None:
        self._log.debug(f'Loading Detailed Occupant data for {self.parent.id}')
        
        self._load_safe(source_data=self._detailed_data, data_key='isActive', target='active')
        self._load_safe(source_data=self._detailed_data, data_key='distance', target='distance_traveled')
        self._load_safe(source_data=self._detailed_data, data_key='location', target='location')
        self._load_safe(source_data=self._detailed_data, data_key='terrainType', target='terrain')
        self._load_safe(source_data=self._detailed_data, data_key='trigger', target='trigger')
        self._load_safe(source_data=self._detailed_data, data_key='tagsApplied', target='_tags')
        self._load_safe(source_data=self._detailed_data, data_key='carbonDioxideVolumeFraction', target='co2_volume_fraction')
        self._load_safe(source_data=self._detailed_data, data_key='carbonMonoxideVolumeFraction', target='co_volume_fraction')
        self._load_safe(source_data=self._detailed_data, data_key='oxygenVolumeFraction', target='o2_volume_fraction')
        self._load_safe(source_data=self._detailed_data, data_key='temperature', target='temperature')
        self._load_safe(source_data=self._detailed_data, data_key='sootVisibility', target='soot_visibility')
        self._load_safe(source_data=self._detailed_data, data_key='*fedTotal*', target='fed_total')
        
        
        try:
            position: Dict = self._detailed_data['position']
            self.position = Point(position['x'], position['y'], position['z'])
        except:
            self.position = Point()
        
        try:
            velocity: Dict = self._detailed_data['velocity']
            
            self.velocity = Vector(velocity['x'], velocity['y'], velocity['z'], velocity['magnitude'])
        except:
            self.velocity = Vector()
        
        try:
            self.trigger = self.parent.parent.triggers[self.trigger]
        except: 
            self.trigger = None
            
        try:
            for tag in self._tags_raw:
                self._tags.append(self.parent.tags[tag])
            self.tags = TagCollection(self._tags)
        except:
            self.tags = TagCollection([])
            
        try:
            self.target = self.parent.parent.targets[self.target]
        except: 
            self.target = None        

class SocialDistanceData(ADataObject):
    """Represents the Social Distance data for a single Occupant.

    Args:
        sd_accum (Dict): The raw parsed data for this Occupant from the sd_accumulated output file 
        sd_transient (Dict): The raw parted data for this Occupant from the sd_transient output file
        parent (Occupant): The :class:`~theng.data.pathfinder.Occupant` to which this data applies.
        
    Attributes:
        detection_radius (float): Distance used to calculate reference distance exposure.
        group_id (Optional[int]): Unique ID for occupant group, if present
        longest_occupant_id (Optional[int]): ID of occupant of greatest exposure (longest within distance SD)
        longest_occupant_name (Optional[str]): Name of occupant of greatest exposure (longest within distance SD)
        longest_occupant_time (Optional[float]): Time within distance SD of occupant of greatest exposure
        occupants_over_60 (List[Occupant]): List of Occupants within detection_radius for > 60 seconds.
        occupants_over_300 (List[Occupant]): List of Occupants within detection_radius for > 300 seconds.
        history (Optional[TimeSeries[SocialDistanceFrame]]): Optional TimeSeries object containing detailes transient SD data.
    """
    
    def __init__(self, *, sd_accum: Dict, sd_transient: Dict, parent: Occupant) -> None:
        super().__init__(parent)
        self._accum_data: Dict = sd_accum
        self._transient_data: Dict = sd_transient
        
        
        #Accumulated Data
        #TODO Verify this typing with a more detailed model
        self.detection_radius: float = float('inf')
        self.group_id: Optional[int] = None
        self.longest_occupant_id: Optional[int] = None
        self.longest_exposure: Optional[Dict] = None
        self._longest_exposure: Optional[Dict] = None;
        self.occupants_over_60: List[Occupant] = []
        self._occupants_over_60: Optional[Dict] = None
        self.occupants_over_300: List[Occupant] = []
        self._occupants_over_300: Optional[Dict] = None
        
        #Transient Data
        self.history: Optional[TimeSeries[SocialDistanceFrame]] = None
        
        self._load()
    
    def _load(self):
        if self._accum_data != None:
            self._load_safe(source_data=self._accum_data, data_key="detectionRadius", target="detection_radius")
            self._load_safe(source_data=self._accum_data, data_key="groupId", target="group_id")
            self._load_safe(source_data=self._accum_data, data_key="longestExposure", target="_longest_exposure")
            self._load_safe(source_data=self._accum_data, data_key="exposureOver60", target="_occupants_over_60")
            self._load_safe(source_data=self._accum_data, data_key="exposureOver300", target="_occupants_over_300")
        
        if self._transient_data != None:
            self.history = TimeSeries(
                time=[_ for _ in self._transient_data.keys()], 
                values=[SocialDistanceFrame(frame_data=_, parent=self.parent) for _ in self._transient_data.values()]
            )
            
    def _update_occupant_references(self):
        if self.history != None and type(self.history) == TimeSeries:
            for frame in self.history.values:
                frame._update_occupant_references()
                
        if self._occupants_over_60 != None:
            for occId in self._occupants_over_60["ids"]:
                occ: Occupant = self.parent.parent.occupants.get_by_id(int(occId))
                self.occupants_over_60.append(occ)
                
        if self._occupants_over_300 != None:
            for occId in self._occupants_over_300["ids"]:
                occ: Occupant = self.parent.parent.occupants.get_by_id(int(occId))
                self.occupants_over_300.append(occ)
                
        if self._longest_exposure != None:
            longestOcc = self.parent.parent.occupants.get_by_id(int(self._longest_exposure["id"]))
            time = self._longest_exposure["time"]
            self.longest_exposure = {
                "occupant": longestOcc,
                "time": time
            }
        
class SocialDistanceFrame(ADataObject):
    """Represents transient SD data at a single time step.

    Args:
        frame_data (Dict): An individual timestep frame of transient Social Distancing data.
        parent (Occupant): The :class:`~theng.data.pathfinder.Occupant` to which this frame of data applies.
        
    Attributes:
        group_id (Optional[int]): Unique ID for occupant group, if present
        closest_occupant (Optional[Dict]): Dictionary containing information about the closest Occupant to this Occupant.
                                           The "occupant" key will yield the Occupant object for the nearest Occupant.
                                           The "dist" key will yield the numerical distance as a float to the nearest Occupant.
        occupants_within_dist (Dict): Dictionary of Occupants within a given distance. Keys are '1.0', '2.0', and '3.0'.
    """
    
    def __init__(self, *, frame_data: Dict, parent: Occupant):
        super().__init__(parent=parent)
        self._data: Dict = frame_data
        
        self.group_id: Optional[int] = None
        self.closest_occupant: Optional[Dict] = None
        self._closest: Optional[Dict] = None
        self.occupants_within_dist: Dict = {}
        self._occs_within: Optional[Dict] = None
                
        self._load()
     
    def _load(self):
        self._load_safe(source_data=self._data, data_key="groupId", target="group_id")
        self._load_safe(source_data=self._data, data_key="closest", target="_closest")
        self._load_safe(source_data=self._data, data_key="occsWithin", target="_occs_within")
        
    def _update_occupant_references(self):
        if self._closest != None:
            self.closest_occupant = {
                "occupant": self.parent.parent.occupants.get_by_id(int(self._closest["id"])),
                "dist": float(self._closest["distance"])
            }
            
        if self._occs_within != None:
            for dist, obj in self._occs_within.items():
                distance: float = float(dist)
                self.occupants_within_dist[distance] = []
                occIds: List[float] = obj['ids']
                for occId in occIds:
                    self.occupants_within_dist[distance].append(self.parent.parent.occupants.get_by_id(int(occId)))