"""
 Copyright 2022 Johns Hopkins University  (Author: Jesus Villalba)
 Apache 2.0  (http://www.apache.org/licenses/LICENSE-2.0)
"""

from typing import List, Optional, Union

import numpy as np
import pandas as pd

from .info_table import InfoTable


class SegmentSet(InfoTable):
    """Class to store information about a speech segment
    Internally, it uses a pandas table.
    """

    def __init__(self, df):
        super().__init__(df)
        if "start" in df and "recording" not in df:
            df["recording"] = df["id"]
            df.fillna(value={"start": 0.0}, inplace=True)

        if "start" not in df and "recording" in df:
            df["start"] = 0.0

        if "recording" in df:
            is_na = df["recording"].isna()
            df.loc[is_na, "recording"] = df.loc[is_na, "id"]

    @property
    def has_time_marks(self):
        return "recording" in self.df and "start" in self.df and "duration" in self.df

    @property
    def has_recording_ids(self):
        return "recording" in self.df

    @property
    def has_recording(self):
        return "recording" in self.df

    def recording(self, ids: Union[np.ndarray, List[str], None] = None):
        if ids is None:
            if "recording" in self.df:
                return self.df["recording"]
            else:
                return self.df["id"]

        if "recording" in self.df:
            return self.df.loc[ids, "recording"]

        return ids

    def image(self, ids: Union[np.ndarray, List[str], None] = None):
        if ids is None:
            if "image" in self.df:
                return self.df["image"]
            else:
                return self.df["id"]

        if "image" in self.df:
            return self.df.loc[ids, "image"]

        return ids

    def video(self, ids: Union[np.ndarray, List[str], None] = None):
        if ids is None:
            if "video" in self.df:
                return self.df["video"]
            else:
                return self.df["video"]

        if "video" in self.df:
            return self.df.loc[ids, "video"]

        return ids

    def recording_ids(self, ids: Union[np.ndarray, List[str], None] = None):
        return self.recording(ids)

    def recording_time_marks(self, ids: Union[np.ndarray, List[str]]):
        if "recording" in self.df:
            recording_name = "recording"
        else:
            recording_name = "id"

        assert "duration" in self.df
        if "start" not in self.df:
            self.df["start"] = 0.0

        return self.df.loc[ids, [recording_name, "start", "duration"]]

    def sample_random_subsegments(
        self,
        subsegments_per_segment: int = 1,
        min_duration: float = 0.0,
        max_duration: Optional[float] = None,
        seg_suffix: Optional[str] = None,
        random_start: bool = True,
        seed: int = 11235813,
        rng: Optional[np.random.Generator] = None,
    ):
        if rng is None:
            rng = np.random.default_rng(seed)

        dfs = []
        for i in range(subsegments_per_segment):
            if max_duration is None:
                duration = rng.uniform(
                    low=min_duration, high=self.df["duration"].values
                )
            else:
                duration = rng.uniform(
                    low=min_duration, high=max_duration, size=(len(self.df),)
                )
                duration = np.minimum(duration, self.df["duration"].values)

            if random_start:
                t_start = rng.uniform(
                    low=0.0, high=self.df["duration"].values - duration
                )
            else:
                t_start = 0.0

            df = self.df.copy()
            df["start"] = t_start
            df["duration"] = duration
            if seg_suffix is None:
                suffix_i = f"-{i}" if subsegments_per_segment > 1 else None
            else:
                suffix_i = (
                    f"{seg_suffix}-{i}" if subsegments_per_segment > 1 else seg_suffix
                )

            if suffix_i is not None:
                df["id"] = df["id"].apply(lambda x: f"{x}-{suffix_i}")

            dfs.append(df)

        df = pd.concat(dfs)
        return SegmentSet(df)