import typing
from pathlib import Path
from warnings import warn

import betterosi
import numpy as np
import pandas as pd
import pyproj
import shapely
from matplotlib import pyplot as plt
from matplotlib.patches import Polygon as PltPolygon

import pandera as pa
import pandera.extensions as extensions

from .asam_odr import MapOdr
from .map import MapOsi, ProjectionOffset

pi_valued = pa.Check.between(-np.pi, np.pi)

@extensions.register_check_method(statistics=['column_name', 'column_value', 'other_column_name', 'other_column_unset_value'])
def other_column_set_on_column_value(pandas_obj, *, column_name:str , column_value, other_column_name: str, other_column_unset_value):
    return ~np.logical_and(pandas_obj[column_name]==column_value, pandas_obj[other_column_name]==other_column_unset_value)

@extensions.register_check_method(statistics=['column_name', 'column_value', 'other_column_name', 'other_column_unset_value'])
def other_column_unset_on_column_value(pandas_obj, *, column_name, column_value, other_column_name: str, other_column_unset_value: int):
    return ~np.logical_and(pandas_obj[column_name]!=column_value, pandas_obj[other_column_name]!=other_column_unset_value)

recording_moving_object_schema = pa.DataFrameSchema(
    title='DataFrame Schema for ASAM OSI GroundTruth of MovingObjects',
    description='some description',
    columns={
        'x': pa.Column(float, 
                       title='MovingObject.base.position.x',
                       description='osi3.Vector3d.x'), 
        'y': pa.Column(float, 
                       title='MovingObject.base.position.y',
                       description='osi3.Vector3d.y'),
        'z': pa.Column(float, 
                       title='MovingObject.base.position.z',
                       description='osi3.Vector3d.z'),
        'vel_x': pa.Column(float, 
                           title='MovingObject.base.velocity.x',
                           description='osi3.Vector3d.x'),
        'vel_y': pa.Column(float, 
                           title='MovingObject.base.velocity.y',
                           description='osi3.Vector3d.y'),
        'vel_z': pa.Column(float, 
                           title='MovingObject.base.velocity.z',
                           description='osi3.Vector3d.z'),
        'acc_x': pa.Column(float, 
                           title='MovingObject.base.acceleration.x',
                           description='osi3.Vector3d.x'),
        'acc_y': pa.Column(float, 
                           title='MovingObject.base.acceleration.y',
                           description='osi3.Vector3d.y'),
        'acc_z': pa.Column(float, 
                           title='MovingObject.base.acceleration.z',
                           description='osi3.Vector3d.z'),
        'length': pa.Column(float, 
                            pa.Check.gt(0), 
                            title='MovingObject.base.dimesion.length',
                            description='osi3.Dimenstion3d.length'),
        'width': pa.Column(float, 
                           pa.Check.gt(0), 
                           title='MovingObject.base.dimesion.width',
                           description='osi3.Dimenstion3d.width'),
        'height': pa.Column(float, 
                            pa.Check.ge(0), 
                            title='MovingObject.base.dimesion.height',
                            description='osi3.Dimenstion3d.height'),
        'type': pa.Column(int, 
                          pa.Check.between(0,4, error=f"Type must be one of {({o.name: o.value for o in betterosi.MovingObjectType})}"), 
                          title="MovingObject.type",
                          description="osi3.MovingObject.Type"),
        'role': pa.Column(int, 
                          pa.Check.between(-1,10, error=f"Type must be one of {({o.name: o.value for o in betterosi.MovingObjectVehicleClassificationRole})}"),
                          title='MovingObject.vehicle_classification.role',
                          description='osi3.MovingObject.VehicleClassification.Role'),
        'subtype': pa.Column(int,  
                             pa.Check.between(-1, 17, error=f"Subtype must be one of {({o.name: o.value for o in betterosi.MovingObjectVehicleClassificationType})}"),
                             title='MovingObject.vehicle_classification.type',
                             description='osi3.MovingObject.VehicleClassification.Type'),
        'roll': pa.Column(float, pi_valued, 
                          title='MovingObject.base.orientation.roll',
                          description='osi3.Orientation3d.roll'),
        'pitch': pa.Column(float, 
                           pi_valued, 
                           title='MovingObject.base.orientation.pitch',
                           description='osi3.Orientation3d.pitch'),
        'yaw': pa.Column(float, 
                         pi_valued, 
                         title='MovingObject.base.orientation.yaw',
                         description='osi3.Orientation3d.yaw'),
        'idx': pa.Column(int, 
                         pa.Check.ge(0), 
                         title='MovingObject.id.value',
                         description='osi3.Identifier.value'),
        'total_nanos': pa.Column(int, 
                                 pa.Check.ge(0), 
                                 title="GroundTruth.timestamp.nanos+1e9*GroundTruth.timestamp.seconds",
                                 description='osi3.Timestamp.nanos, osi3.Timestamp.seconds'),
    },
    unique=['idx','total_nanos'],
    checks=[
        pa.Check.other_column_set_on_column_value(
            'type', 
            int(betterosi.MovingObjectType.TYPE_VEHICLE), 
            'role', 
            -1, 
            error="`role` is `-1` despite type beeing `TYPE_VEHICLE`"),
        pa.Check.other_column_unset_on_column_value(
            'type', 
            int(betterosi.MovingObjectType.TYPE_VEHICLE), 
            'role', 
            -1, 
            error="`role` is set despite type not beeing `TYPE_VEHICLE`"),
        pa.Check.other_column_set_on_column_value(
            'type', 
            int(betterosi.MovingObjectType.TYPE_VEHICLE), 
            'subtype', 
            -1, 
            error="`subtype` is `-1` despite type beeing `TYPE_VEHICLE`"),
        pa.Check.other_column_unset_on_column_value(
            'type', 
            int(betterosi.MovingObjectType.TYPE_VEHICLE), 
            'subtype', 
            -1, 
            error="`subtype` is set despite type not beeing `TYPE_VEHICLE`"),
    ])


#recording_internal_mv_schema = recording_moving_object_schema.add_columns({
#    'polygon': pa.Column(shapely.Polygon),
#    'frame': pa.Column(int, pa.Check.ge(0))
#})

def timestamp2ts(timestamp: betterosi
.Timestamp):
    return timestamp.seconds*1_000_000_000+timestamp.nanos

def nearest_interp(xi, x, y):
    # https://stackoverflow.com/a/21003629
    idx = np.abs(x - xi[:,None])
    return y[idx.argmin(axis=1)]
        

class MovingObject:
    def __init__(self, recording, idx):
        super().__init__()
        self.idx = int(idx)
        self._recording = recording
       
        self._df = self._recording._df.loc[self._recording._df['idx']==self.idx]
        
        for k in ['x', 'y', 'z', 'vel_x', 'vel_y', 'vel_z', 'acc_x', 'acc_y', 'acc_z', 'yaw', 'roll', 'pitch', 'polygon']:
            setattr(self, k, self._df.loc[:, k].values)
        self.vel = np.linalg.norm([self.vel_x, self.vel_y], axis=0)
        self.timestamps = self._df.loc[:,'total_nanos'].values/1e9
        for k in ['length', 'width', 'height']:
            setattr(self, f'{k}s', self._df.loc[:, k].values)
            setattr(self, k, np.mean(self._df.loc[:, k].values))

        self.type = betterosi.MovingObjectType(self._df.loc[:,'type'].iloc[0])
        subtype_int = self._df.loc[:,'subtype'].iloc[0]
        self.subtype = betterosi.MovingObjectVehicleClassificationType(subtype_int) if subtype_int != -1 else None
        role_int = self._df.loc[:,'role'].iloc[0]
        self.role = betterosi.MovingObjectVehicleClassificationRole(role_int) if role_int != -1 else None
        self.birth = int(self._df.loc[:,'frame'].iloc[0])
        self.end = int(self._df.loc[:,'frame'].iloc[-1])   
    
    def set(self, k, val):
        self._recording._df.loc[self._recording._df['idx']==self.idx, k] = val

    @property
    def nanos(self):
        return self.timestamps*1e9

    def plot(self, ax: plt.Axes):
        ax.plot(self.x, self.y, label=str(self.idx), c='red', alpha=.5)
        pass
    
    def plot_mv_frame(self, ax: plt.Axes, frame: int):
        polys = self._df[self._df['frame']==frame]['polygon'].values
        for p in polys:
            ax.add_patch(PltPolygon(p.exterior.coords, fc='red', alpha=0.2))
    


class Recording:
    _MovingObjectClass: typing.ClassVar = MovingObject
    
    @staticmethod
    def _get_polygons(df):
        c2f = df['length']/2
        c2l = df['width']/2
        x = df['x']
        y = df['y']
        cosyaw = np.cos(df['yaw'])
        sinyaw = np.sin(df['yaw']) 
        polys = np.array([
            ((x + (+c2f) * cosyaw- (+c2l) * sinyaw),(y + (+c2f) * sinyaw + (+c2l) * cosyaw)),
            ((x + (+c2f) * cosyaw- (-c2l) * sinyaw),(y + (+c2f) * sinyaw + (-c2l) * cosyaw)),
            ((x + (-c2f) * cosyaw- (-c2l) * sinyaw),(y + (-c2f) * sinyaw + (-c2l) * cosyaw)),
            ((x + (-c2f) * cosyaw- (+c2l) * sinyaw),(y + (-c2f) * sinyaw + (+c2l) * cosyaw))
        ]).swapaxes(0,2).swapaxes(1,2)
        return shapely.polygons(polys)
    
    @staticmethod
    def get_moving_object_ground_truth(nanos: int, df: pd.DataFrame, host_vehicle=None, validate=True) -> betterosi.GroundTruth:
        if validate:
            recording_moving_object_schema.validate(df, lazy=True)
        def get_object(row):
            return betterosi.MovingObject(
                id=betterosi.Identifier(value=row['idx']),
                type=betterosi.MovingObjectType(row['type']),
                base=betterosi.BaseMoving(
                    dimension=betterosi.Dimension3D(length=row['length'], width=row['width'], height=row['width']),
                    position=betterosi.Vector3D(x=row['x'], y=row['y'], z=row['z']),
                    orientation=betterosi.Orientation3D(roll=row['roll'], pitch=row['pitch'], yaw=row['yaw']),
                    velocity=betterosi.Vector3D(x=row['vel_x'], y=row['vel_y'], z=row['vel_z']),
                    acceleration=betterosi.Vector3D(x=row['acc_x'], y=row['acc_y'], z=row['acc_z']),
                ),
                vehicle_classification=betterosi.MovingObjectVehicleClassification(type=row['subtype'], role=row['role'])
            )
        mvs = list(df.apply(get_object, axis=1).values)
        gt = betterosi.GroundTruth(
            version=betterosi.InterfaceVersion(version_major=3, version_minor=7, version_patch=9),
            timestamp=betterosi.Timestamp(seconds=int(nanos//1_000_000_000), nanos=int(nanos%1_000_000_000)),
            host_vehicle_id=betterosi.Identifier(value=0) if host_vehicle is None else betterosi.Identifier(value=host_vehicle),
            moving_object=mvs
        )
        return gt

    def __init__(self, df, map=None, projections=None, host_vehicle=None, validate=True):
        if validate:
            recording_moving_object_schema.validate(df, lazy=True)
        super().__init__()
        self.nanos2frame = {n: i for i, n in enumerate(df.total_nanos.unique())}
        df['frame'] = df.total_nanos.map(self.nanos2frame)
        if 'polygon' not in df.columns:
            df['polygon'] = self._get_polygons(df)
        self.projections = projections
        self._df = df
        self.map = map
        self.moving_objects = {int(idx): self._MovingObjectClass(self, idx) for idx in self._df['idx'].unique()}
        self.host_vehicle = host_vehicle

    def to_osi_gts(self) -> list[betterosi.GroundTruth]:
        gts = [self.get_moving_object_ground_truth(nanos, group_df, host_vehicle=self.host_vehicle, validate=False) for nanos, group_df in self._df.groupby('total_nanos')]
        
        if self.map is not None and isinstance(self.map, MapOsi):
            gts[0].lane_boundary = [b._osi for b in self.map.lane_boundaries.values()]
            gts[0].lane = [l._osi for l in self.map.lanes.values()]
        return gts
        
    @classmethod
    def from_osi_gts(cls, gts: list[betterosi.GroundTruth], validate: bool = True):
        mvs = []
        map = None
        projs = []
        for i, gt in enumerate(gts):
            total_nanos = gt.timestamp.seconds*1_000_000_000 + gt.timestamp.nanos
            if gt.proj_frame_offset is not None and gt.proj_frame_offset.position is None:
                    raise ValueError(f'Offset of {i}th ground truth message (total_nanos={total_nanos}) is set without position.')
            projs.append(dict(
                proj_string = gt.proj_string,
                projection = pyproj.CRS.from_proj4(gt.proj_string) if gt.proj_string is not None and gt.proj_string != "" else None,
                offset = ProjectionOffset(
                    x=gt.proj_frame_offset.posiiton.x,
                    y=gt.proj_frame_offset.position.y,
                    z=gt.proj_frame_offset.position.z,
                    yaw=gt.proj_frame_offset.yaw) if gt.proj_frame_offset is not None else None
            ))
            mvs += [dict(
                total_nanos = total_nanos,
                idx = mv.id.value,
                x = mv.base.position.x,
                y = mv.base.position.y,
                z = mv.base.position.z,
                vel_x = mv.base.velocity.x,
                vel_y = mv.base.velocity.y,
                vel_z = mv.base.velocity.z,
                acc_x = mv.base.acceleration.x,
                acc_y = mv.base.acceleration.y,
                acc_z = mv.base.acceleration.z,
                length = mv.base.dimension.length,
                width = mv.base.dimension.width,
                height = mv.base.dimension.height,
                roll = mv.base.orientation.roll,
                pitch = mv.base.orientation.pitch,
                yaw = mv.base.orientation.yaw,
                type = mv.type,
                role = mv.vehicle_classification.role if mv.type == betterosi.MovingObjectType.TYPE_VEHICLE else -1,
                subtype = mv.vehicle_classification.type if mv.type == betterosi.MovingObjectType.TYPE_VEHICLE else -1
            ) for mv in gt.moving_object]
            
            if map is None:
                map = MapOsi.create(gt)
        df_mv = pd.DataFrame(mvs).sort_values(by=['total_nanos', 'idx']).reset_index(drop='index')
        return cls(df_mv, map, projections=projs, validate=validate)
    
    @classmethod
    def from_file(cls, filepath, xodr_path: str|None = None, validate:bool = True, skip_odr_parse: bool = False):
        gts = betterosi.read(filepath, return_ground_truth=True,)
        r = cls.from_osi_gts(gts, validate=validate)
        if xodr_path is not None:
            r.map = MapOdr.from_file(xodr_path, skip_parse=skip_odr_parse)
        elif Path(filepath).suffix==".mcap":
            try:
                r.map = MapOdr.from_file(filepath, skip_parse=skip_odr_parse)
            except StopIteration:
                if r.map is not None:
                    warn("No map provided in mcap. OSI GroundTruth map is used!")
        if r.map is None:
            warn("No xodr map provided in mcap nor OSI map in GroundTruth!")
        return r
    
    def to_mcap(self, filepath):
        if Path(filepath).suffix != '.mcap':
            raise ValueError()
        gts = self.to_osi_gts()
        with betterosi.Writer(filepath) as w:
            for gt in gts:
                w.add(gt)
            if isinstance(self.map, MapOdr):
                w.add(self.map.to_osi(), topic='ground_truth_map', log_time=0)

        
    def to_hdf(self, filename, key='moving_object'):
        #!pip install tables
        self._df.drop(columns=['polygon','frame']).to_hdf(filename, key=key)
        
    @classmethod
    def from_hdf(cls, filename, key='moving_object'):
        df = pd.read_hdf(filename, key=key)
        return cls(df, map=None, host_vehicle=None)
        
    def interpolate(self, new_nanos: list[int]|None = None, hz: float|None = None):
        df = self._df
        if new_nanos is None and hz is None:
            new_nanos = np.linspace(df.total_nanos.min(), df.total_nanos.max(), df.frame.max()-df.frame.min(), dtype=int)
        elif hz is not None:
            step = 1_000_000_000*hz
            new_nanos = np.arange(start=df.total_nanos.min(), stop=df.total_nanos.max()+1, step=step, dtype=int)
        else:
            new_nanos = np.array(new_nanos)
        new_dfs = []
        for idx, track_df in df.groupby('idx'):
            track_data = {}
            track_new_nanos = new_nanos[track_df.frame.min()-df.frame.min():track_df.frame.max()-df.frame.min()+1]
            for c in ['x', 'y', 'z', 'vel_x', 'vel_y', 'vel_z', 'acc_x',
                'acc_y', 'acc_z', 'length', 'width', 'height']:
                track_data[c] = np.interp(track_new_nanos, track_df['total_nanos'], track_df[c])
            for c in ['type', 'subtype', 'role']:
                track_data[c] = nearest_interp(track_new_nanos, track_df['total_nanos'].values, track_df[c].values)
            for c in ['roll', 'pitch', 'yaw']:
                track_data[c] = np.mod(np.interp(track_new_nanos, track_df['total_nanos'], np.unwrap(track_df[c], period=np.pi)), np.pi)
            new_track_df = pd.DataFrame(track_data)
            new_track_df['idx'] = idx
            new_track_df['total_nanos'] = track_new_nanos
            new_dfs.append(new_track_df)
        new_df = pd.concat(new_dfs)
        return self.__init__(new_df, self.map, self.host_vehicle)

    def plot(self, ax = None, legend=False) -> plt.Axes:
        if ax is None:
            fig, ax = plt.subplots(1,1)
        if self.map:
            self.map.plot(ax)
        for ru in self.moving_objects.values():
            ru.plot(ax)
        if legend:
            ax.legend(loc='center left', bbox_to_anchor=(1, 0.5))
        return ax   
    
    def plot_frame(self, frame: int, ax = None):
        ax = self.plot(ax=ax)
        self.plot_mv_frame(ax, frame=frame)
        return ax
        
    def plot_mv_frame(self, ax: plt.Axes, frame: int):
        polys = self._df[self._df['frame']==frame]['polygon'].values
        for p in polys:
            ax.add_patch(PltPolygon(p.exterior.coords, fc='red'))
     
