import os
from datetime import datetime

import pandas as pd

from srdatasets.datasets.dataset import Dataset
from srdatasets.datasets.utils import download_url, extract


class Lastfm1K(Dataset):

    __url__ = "http://mtg.upf.edu/static/datasets/last.fm/lastfm-dataset-1K.tar.gz"
    __corefile__ = os.path.join("lastfm-dataset-1K", "userid-timestamp-artid-artname-traid-traname.tsv")

    def download(self):
        download_url(self.__url__, self.rawpath)
        extract(self.rawpath, self.rootdir)

    def transform(self, item_type):
        """ item_type can be `artist` or `song`
        """
        df = pd.read_csv(
            self.rootdir.joinpath(self.__corefile__),
            sep="\t",
            names=["user_id", "timestamp", "artist_id", "artist_name", "song_id", "song_name"],
            usecols=[0, 1, 2, 4],
            converters={"timestamp": lambda x: int(datetime.strptime(x, "%Y-%m-%dT%H:%M:%SZ").timestamp())})
        if item_type == "song":
            df = df.drop("artist_id", axis=1).rename(columns={"song_id": "item_id"})
        else:
            df = df.drop("song_id", axis=1).rename(columns={"artist_id": "item_id"})
        return df
