import datetime
from argparse import ArgumentDefaultsHelpFormatter, ArgumentParser
from logging import INFO, Logger, getLogger
from random import choice, randint, random
from typing import Any, List

import pandas as pd
import polars as pl
from pyais.ais_types import AISType

from damast.core.metadata import DataCategory
from damast.core.types import DataFrame, XDataFrame
from damast.domains.maritime.ais import vessel_types
from damast.domains.maritime.ais.navigational_status import AISNavigationalStatus
from damast.domains.maritime.data_specification import (
    MMSI,
    ColumnName,
    CourseOverGround,
    SpeedOverGround,
    )

_log: Logger = getLogger(__name__)
_log.setLevel(INFO)

DATETIME_FORMAT = "%Y-%m-%d %H:%M:%S"
AIS_DATA_COLUMNS = [
    "mmsi",
    "lon",
    "lat",
    "date_time_utc",
    "sog",
    "cog",
    "true_heading",
    "nav_status",
    "rot",
    "message_nr",
    "source"
]

TRAJECTORY_MIN_SIZE = 10
TRAJECTORY_MAX_SIZE = 200

AISTestDataSpec = {
    "annotations": {"comment": "This is a autogenerated test data set"},
    "columns": [
        {"name": "mmsi", "is_optional": False, "category": DataCategory.STATIC},
        # , "value_range": CyclicMinMax(-180.0, 180.0)},
        {"name": "lon", "is_optional": False, "unit": "deg", "category": DataCategory.DYNAMIC},
        # , "value_range": CyclicMinMax(-90.0, 90.0)},
        {"name": "lat", "is_optional": False, "unit": "deg", "category": DataCategory.DYNAMIC},
        {"name": "date_time_utc", "is_optional": False, "category": DataCategory.DYNAMIC,
         "representation_type": "str"},
        {"name": "sog", "is_optional": False, "category": DataCategory.DYNAMIC},
        {"name": "cog", "is_optional": False, "category": DataCategory.DYNAMIC},
        {"name": "true_heading", "is_optional": False, "category": DataCategory.DYNAMIC},
        {"name": "nav_status", "is_optional": False, "category": DataCategory.DYNAMIC},
        {"name": "rot", "is_optional": False, "category": DataCategory.DYNAMIC},
        {"name": "message_nr", "is_optional": False, "category": DataCategory.DYNAMIC},
        {"name": "source", "is_optional": False, "category": DataCategory.DYNAMIC},
    ]
}


class AISTestData:
    dataframe: DataFrame = None

    def __init__(self,
                 number_of_trajectories: int,
                 min_length: int = TRAJECTORY_MIN_SIZE,
                 max_length: int = TRAJECTORY_MAX_SIZE):

        self.number_of_trajectories = number_of_trajectories
        self.min_length = min_length
        self.max_length = max_length

        self.dataframe: DataFrame = self._generate_data()

    @staticmethod
    def generate_trajectory(min_size: int, max_size: int) -> List[List[Any]]:
        """Generate a trajectory for a single vessel

        :param min_size: Minimum size of trajectory
        :param max_size: Maximum size of trajectory
        :return: A nested list describing the trajectory
        """
        # Create some invalid messages as well
        range_ext = 0.05
        mmsi = randint(int((1-range_ext)*MMSI.min_value), int((1+range_ext)*MMSI.max_value))
        lat_start = (random() * 180.0) - 90.0
        lon_start = (random() * 360.0) - 180.0

        last_day = datetime.datetime(year=2022, month=12, day=31)
        start_time = randint(0, int(last_day.timestamp()))

        sog_start = random() * SpeedOverGround.min_value / SpeedOverGround.max_value
        cog_start = random() * CourseOverGround.min_value / CourseOverGround.max_value
        heading_start = cog_start

        rot_start = 0.0

        # Choose a random message ID from a subset of available types
        msg_id = choice([AISType.POS_CLASS_A1, AISType.POS_CLASS_A2, AISType.POS_CLASS_A3])

        trajectory_length = randint(min_size, max_size)
        trajectory = []
        for _ in range(0, trajectory_length):
            lat_start += random() * 0.1 - 0.05
            lon_start += random() * 0.1 - 0.05

            # message within 2 min
            start_time += random() * 120.0

            sog_start += random() * 4 - 2
            cog_start += random() * 1 - 0.5
            heading_start = cog_start + random() / 10.0

            nav_status = choice([
                AISNavigationalStatus.EngagedInFishing.value,
                AISNavigationalStatus.UnderWayUsingEngine.value,
                AISNavigationalStatus.AtAnchor.value
            ])
            # 5 percent of all messages should be missing a time-stamp
            q = random()
            if q > 0.1:
                timestamp = datetime.datetime.fromtimestamp(start_time).strftime(
                    DATETIME_FORMAT)    # 2022-09-01 12:55:35
            else:
                timestamp = None
            trajectory.append(
                [mmsi,
                 lon_start,
                 lat_start,
                 timestamp,
                 sog_start,
                 cog_start,
                 heading_start,
                 nav_status,
                 rot_start,
                 msg_id,
                 choice(["s", "g"])  # Create some ground data that should be removed in processing
                 ])
        return trajectory

    def _generate_data(self) ->DataFrame:
        df = None
        for i in range(0, self.number_of_trajectories):
            trajectory = AISTestData.generate_trajectory(min_size=self.min_length, max_size=self.max_length)
            t_df = pd.DataFrame(trajectory, columns=AIS_DATA_COLUMNS)
            if df is None:
                df = t_df
            else:
                df = pd.concat([df, t_df], axis=0, ignore_index=True)
        # shuffle rows
        df = df.sample(frac=1, axis="index").reset_index(drop=True)
        return XDataFrame.from_pandas(df)

    def generate_vessel_type_data(self) -> DataFrame:
        """Generate a :class:`DataFrame` with data imitating vessel-type info

        Returns:
            A dataframe with an `ColumnName.MMSI` column and a `ColumnName.VESSEL_TYPE` column.
        """
        vessel_type_data = []
        for mmsi in self.dataframe[ColumnName.MMSI.lower()].unique():
            vessel_type_data.append([mmsi, vessel_types.Fishing.to_id()])
        df = pd.DataFrame(vessel_type_data, columns=[ColumnName.MMSI, ColumnName.VESSEL_TYPE])
        return XDataFrame.from_pandas(df)

    def generate_fishing_vessel_type_data(self) -> DataFrame:
        """Generate a `vaex.Dataframe` with data imitating vessel info from Global Fishing Watch

        Returns:
            A dataframe with an `ColumnName.MMSI` (lower-cased) column and a `ColumnName.VESSEL_TYPE` column.
        """
        fishing_vessel_type_data = []
        for mmsi in self.dataframe[ColumnName.MMSI.lower()].unique():
            fishing_vessel_type_data.append([mmsi, choice([
                vessel_types.DriftingLonglines.to_id(),
                vessel_types.PoleAndLine.to_id(),
                vessel_types.PotsAndTraps.to_id()
            ])])

        df = pd.DataFrame(fishing_vessel_type_data, columns=[ColumnName.MMSI.lower(), ColumnName.VESSEL_TYPE_GFW])

        return XDataFrame.from_pandas(df)

    def generate_anchorage_type_data(self) -> DataFrame:
        # Generate anchorages close to starting points of vessel trajectories
        mmsi = ColumnName.MMSI.lower()
        first = self.dataframe.group_by(mmsi).first()
        first = first.rename({ "lat": "latitude", "lon": "longitude"})
        df = first.with_columns(
                pl.col("latitude")*1.05,
                pl.col("longitude")*-0.95,
            )
        return df


if __name__ == "__main__":
    parser = ArgumentParser(formatter_class=ArgumentDefaultsHelpFormatter,
                            description="A simple generator for ais data that can be used as input to "
                                        + "data processing stage")

    parser.add_argument("-t", "--number-of-trajectories", default=100, type=int)
    parser.add_argument("-o", "--output", default="test-data.hdf5", type=str)

    args = parser.parse_args()

    ais_test_data = AISTestData(number_of_trajectories=args.number_of_trajectories)
    ais_test_data.dataframe.export_hdf5(args.output)
    _log.info("Written: {args.output}")
