import numpy as np

from riskmap import *
from riskmap.BallisticFall import GenerateImpactMap as generate_impact_map
from riskmap.BallisticFall import GenerateBallisticFall as generate_ballistic_fall
from riskmap.BallisticFall import BallisticFall as ballistic_fall


class BallisticRiskMap():
    def __init__(self, ground_map, aircraft, n_sigma, hdg_samples, num_falls, sim_time, v0, altitude_slices, hdg_sigma):
        self.hdg_samples = np.arange(0, np.pi * 2, 2 * np.pi / hdg_samples)
        self.ground_map = ground_map
        self.aircraft = aircraft
        self.n_sigma = n_sigma
        self.risk_map = dict()
        self.altitude_slices = altitude_slices
        for hdg in self.hdg_samples:
            if hdg not in self.risk_map:
                self.risk_map[np.round(hdg, 4)] = dict()

            impact_maps = generate_impact_map(num_falls, sim_time, v0, hdg, hdg_sigma, altitude_slices, aircraft)
            for impact_map in impact_maps:
                self.risk_map[np.round(hdg, 4)][np.round(impact_map.altitude, 0)] = impact_map.generate_probability_map(n_sigma, ground_map)
        dummy_fall = generate_ballistic_fall(aircraft.c, aircraft.S, aircraft.m, 0., v0, sim_time)
        
        # Ballistic fall trajectory is solved numerically, and so it is sampled in time.
        # Too many samples are then above the same cell in the ground map
        # Reduce the number of samples by resampling in distance
        dist_samples = np.arange(0, dummy_fall.dist[-1], int(np.round(self.ground_map.resolution / 2, 0)))
        dist_samples = np.append(dist_samples, dummy_fall.dist[-1])
        idxs = [np.where(np.array(dummy_fall.dist) <= dist_sample)[0][-1] for dist_sample in dist_samples]
        
        dist = dummy_fall.dist[idxs]
        alt = dummy_fall.alt[idxs]
        v_h = dummy_fall.v_h[idxs]
        v_v = dummy_fall.v_v[idxs]
        self.dummy_fall = ballistic_fall(dist, alt, v_h, v_v, dummy_fall.init_hdg)

    """
    Method to determine a risk induced by potential ballistic failure along the given trajectory.
    Args:   
        Path path   # path to be assesed
    Returns:
        float risk  # risk induced by potential ballistic fall along the path
    """
    def get_risk(self, path):

        # TODO -- implement me
        risk = 0
            
        return risk

        


    

           
