import json
import pickle
from collections import defaultdict
from dataclasses import dataclass
from pathlib import Path

import pandas as pd
import torch
from omegaconf import MISSING

from s3prl.dataio.dataset import EncodeMultiLabel, LoadAudio
from s3prl.dataio.encoder import CategoryEncoder
from s3prl.dataio.sampler import FixedBatchSizeBatchSampler
from s3prl.nn.hear import HearFullyConnectedPrediction
from s3prl.task.scene_prediction import ScenePredictionTask

from ._hear_util import resample_hear_corpus
from .superb_sid import SuperbSID

__all__ = ["HearFSD"]


def hear_scene_trainvaltest(
    target_dir: str,
    cache_dir: str,
    dataset_root: str,
    get_path_only: bool = False,
):
    target_dir = Path(target_dir)

    resample_hear_corpus(dataset_root, target_sr=16000)

    dataset_root = Path(dataset_root)
    wav_root: Path = dataset_root / "16000"

    train_csv = target_dir / "train.csv"
    valid_csv = target_dir / "valid.csv"
    test_csv = target_dir / "test_csv"

    if get_path_only:
        return train_csv, valid_csv, [test_csv]

    def load_json(filepath):
        with open(filepath, "r") as fp:
            return json.load(fp)

    def split_to_df(split: str) -> pd.DataFrame:
        meta = load_json(dataset_root / f"{split}.json")
        data = defaultdict(list)
        for k in list(meta.keys()):
            data["id"].append(k)
            data["wav_path"].append(wav_root / split / k)
            data["labels"].append(" ; ".join([str(label).strip() for label in meta[k]]))
        return pd.DataFrame(data=data)

    split_to_df("train").to_csv(train_csv, index=False)
    split_to_df("valid").to_csv(valid_csv, index=False)
    split_to_df("test").to_csv(test_csv, index=False)

    return train_csv, valid_csv, [test_csv]


class HearFSD(SuperbSID):
    def default_config(self) -> dict:
        return dict(
            start=0,
            stop=None,
            target_dir=MISSING,
            cache_dir=None,
            remove_all_cache=False,
            prepare_data=dict(
                dataset_root=MISSING,
            ),
            build_batch_sampler=dict(
                train=dict(
                    batch_size=10,
                    shuffle=True,
                ),
                valid=dict(
                    batch_size=1,
                ),
                test=dict(
                    batch_size=1,
                ),
            ),
            build_upstream=dict(
                name=MISSING,
            ),
            build_featurizer=dict(
                layer_selections=None,
                normalize=False,
            ),
            build_downstream=dict(
                hidden_layers=2,
                pooling_type="MeanPooling",
            ),
            build_model=dict(
                upstream_trainable=False,
            ),
            build_task=dict(
                prediction_type="multilabel",
                scores=["mAP", "top1_acc", "d_prime", "aucroc"],
            ),
            build_optimizer=dict(
                name="Adam",
                conf=dict(
                    lr=1.0e-3,
                ),
            ),
            build_scheduler=dict(
                name="ExponentialLR",
                gamma=0.9,
            ),
            save_model=dict(),
            save_task=dict(),
            train=dict(
                total_steps=40000,
                log_step=100,
                eval_step=1000,
                save_step=100,
                gradient_clipping=1.0,
                gradient_accumulate=1,
                valid_metric="mAP",
                valid_higher_better=True,
                auto_resume=True,
                resume_ckpt_dir=None,
            ),
            evaluate=dict(),
        )

    def prepare_data(
        self,
        prepare_data: dict,
        target_dir: str,
        cache_dir: str,
        get_path_only: bool = False,
    ):
        return hear_scene_trainvaltest(
            **self._get_current_arguments(flatten_dict="prepare_data")
        )

    def build_encoder(
        self,
        build_encoder: dict,
        target_dir: str,
        cache_dir: str,
        train_csv_path: str,
        valid_csv_path: str,
        test_csv_paths: list,
        get_path_only: bool = False,
    ):
        encoder_path = Path(target_dir) / "encoder.pkl"
        if get_path_only:
            return encoder_path

        train_csv = pd.read_csv(train_csv_path)
        valid_csv = pd.read_csv(valid_csv_path)
        test_csvs = [pd.read_csv(path) for path in test_csv_paths]
        all_csv = pd.concat([train_csv, valid_csv, *test_csvs])
        all_labels = []
        for rowid, row in all_csv.iterrows():
            labels = str(row["labels"]).split(";")
            labels = [l.strip() for l in labels]
            all_labels.extend(labels)

        encoder = CategoryEncoder(all_labels)
        with open(encoder_path, "wb") as f:
            pickle.dump(encoder, f)

        return encoder_path

    def build_dataset(
        self,
        build_dataset: dict,
        target_dir: str,
        cache_dir: str,
        mode: str,
        data_csv: str,
        encoder_path: str,
        frame_shift: int,
    ):
        df = pd.read_csv(data_csv)
        ids = df["id"].tolist()
        wav_paths = df["wav_path"].tolist()
        labels = [
            [single_label.strip() for single_label in str(label_str).split(";")]
            for label_str in df["labels"].tolist()
        ]
        with open(encoder_path, "rb") as f:
            encoder = pickle.load(f)

        audio_loader = LoadAudio(wav_paths)
        label_encoder = EncodeMultiLabel(labels, encoder)

        class Dataset:
            def __len__(self):
                return len(audio_loader)

            def __getitem__(self, index: int):
                audio = audio_loader[index]
                label = label_encoder[index]
                return {
                    "x": audio["wav"],
                    "x_len": audio["wav_len"],
                    "y": label["binary_labels"],
                    "labels": label["labels"],
                    "unique_name": ids[index],
                }

        dataset = Dataset()
        return dataset

    def build_batch_sampler(
        self,
        build_batch_sampler: dict,
        target_dir: str,
        cache_dir: str,
        mode: str,
        data_csv: str,
        dataset,
    ):
        @dataclass
        class Config:
            train: dict = None
            valid: dict = None
            test: dict = None

        conf = Config(**build_batch_sampler)
        return FixedBatchSizeBatchSampler(dataset, **(conf.train or {}))

    def build_downstream(
        self,
        build_downstream: dict,
        downstream_input_size: int,
        downstream_output_size: int,
        downstream_input_stride: int,
    ):
        return HearFullyConnectedPrediction(
            downstream_input_size, downstream_output_size, **build_downstream
        )

    def build_task(
        self,
        build_task: dict,
        model: torch.nn.Module,
        encoder,
        valid_df: pd.DataFrame = None,
        test_df: pd.DataFrame = None,
    ):
        return ScenePredictionTask(model, encoder, **build_task)