# Copyright (C) 2023 Jae-Won Chung <jwnchung@umich.edu>
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""
Generate the Capriccio dataset.

This script takes as input the path to the Sentiment140 dataset
and slices it into 38 overlapping slices. Each slice has 425000
tweets for training (%d_train.json) and 75000 tweets for eval
(%d_val.json).

Usage:
    python capriccio/generate.py /path/to/sentiment140.json
"""


import argparse
import os
import warnings
from datetime import datetime, timedelta
from pathlib import Path

import datasets
import numpy as np
import pandas as pd


def main(output_dir: str) -> None:
    """Run the main routine."""
    # Prepare raw dataset
    print("Preparing raw dataset.")
    df = datasets.load_dataset("sentiment140")["train"].to_pandas()
    with warnings.catch_warnings():
        warnings.filterwarnings("ignore")
        df["date"] = pd.to_datetime(df.date)
    df["ind"] = df.index
    df = df.set_index("date")
    df = df.sort_index()
    df = df.rename({"sentiment": "label"}, axis=1)
    df = df.drop(["user", "query"], axis=1)

    # Slice datasets
    stride = timedelta(days=1)
    size = 500_000

    print(f"Slicing datasets with stride {stride} and size {size}.")

    slice_index = 1
    sliced = []
    # Skip April since there are too many days with no tweets.
    now = datetime(year=2009, month=5, day=1)
    end = df.index.max()
    while now < end:
        loc = df.index.get_loc(now.strftime("%m/%d/%Y")).start
        slice_df = df.iloc[loc : loc + size]
        if len(slice_df) < size:
            break

        # Compute sample overlap ratio
        if sliced:
            overlap = len(sliced[-1].merge(slice_df, how="inner", on="ind"))
            print(
                f"{slice_index:2d}: {slice_df.index.min()} ~ {slice_df.index.max()}, overlap = {overlap/size:.3f}"
            )
        else:
            print(f"{slice_index:2d}: {slice_df.index.min()} ~ {slice_df.index.max()}")

        sliced.append(slice_df)

        slice_index += 1
        now += stride

    print(f"{len(sliced)} datasets of size {size} were created.")

    # Split train and validation sets and save
    print("Sampling validation set and saving.")
    seed = 42
    train_frac = 0.85
    save_dir = Path(output_dir)
    os.makedirs(save_dir, exist_ok=True)

    for slice_index, dataset in enumerate(sliced):
        ind = np.random.default_rng(seed).permutation(len(dataset))
        shuffled = dataset.iloc[ind]
        boundary = int(len(dataset) * train_frac)
        train = shuffled.iloc[:boundary].drop(["ind"], axis=1).reset_index()
        val = shuffled.iloc[boundary:].drop(["ind"], axis=1).reset_index()
        train.to_json(
            save_dir / f"{slice_index+1}_train.json", orient="records", lines=True
        )
        val.to_json(
            save_dir / f"{slice_index+1}_val.json", orient="records", lines=True
        )


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--output-dir", "-o", required=True, help="Directory to save Capriccio"
    )
    args = parser.parse_args()

    main(args.output_dir)
