import os
import re
import shutil
import subprocess
import hashlib
import urllib.request
from glob import glob
from typing import List, Tuple, Optional

import numpy as np
import pandas as pd
from joblib import Parallel, delayed
from tensorflow.keras.preprocessing.sequence import pad_sequences

from useckit.v2.dataset import WindowslicedDataset, Dataset


_DEF_DIRNAME = "chi-2021"
_FILE1 = "Data_Set_for_Understanding_User_Identification_in_Virtual_Reality_Through_Behavioral_Biometrics_and_the_Effect_of_Body_Normalization.7z.001"
_FILE2 = "Data_Set_for_Understanding_User_Identification_in_Virtual_Reality_Through_Behavioral_Biometrics_and_the_Effect_of_Body_Normalization.7z.002"
_URL1 = "https://hci.informatik.uni-due.de/fileadmin/fileupload/I-HCI/Paper/" + _FILE1
_URL2 = "https://hci.informatik.uni-due.de/fileadmin/fileupload/I-HCI/Paper/" + _FILE2
_SHA1 = "1235b224ba131c36fc70a29a2e381f8ec0f5d7aa51c6ad09e576068c2c1c8d7d"
_SHA2 = "3110f9603a07a0259eabff73c399ebb679880f28390d48ca6480623e2c79bda2"
_FILENAME_PATTERN = r'(?P<activity>\w+)_p(?P<pid>\d+)_(?P<condition>\w+)_session(?P<session>\d+)_repetition(?P<rep>\d+)\.csv'


def _download_dataset(dataset_dir: str, reinstall: bool = False) -> str:
    if reinstall and os.path.exists(dataset_dir):
        shutil.rmtree(dataset_dir)
    if not os.path.exists(dataset_dir):
        os.makedirs(dataset_dir, exist_ok=True)
        zip1 = os.path.join(dataset_dir, _FILE1)
        zip2 = os.path.join(dataset_dir, _FILE2)
        # Download
        urllib.request.urlretrieve(_URL1, zip1)
        urllib.request.urlretrieve(_URL2, zip2)
        # Hash check
        def sha256(path):
            with open(path, "rb") as f:
                return hashlib.sha256(f.read()).hexdigest()
        h1, h2 = sha256(zip1), sha256(zip2)
        print(f"sha256sum of {zip1} is {h1}")
        print(f"sha256sum of {zip2} is {h2}")
        assert h1 == _SHA1, "hashsum did not match expectation for .001"
        assert h2 == _SHA2, "hashsum did not match expectation for .002"
        # Extract using 7z
        try:
            subprocess.run(["7z", "x", zip1, "-o" + dataset_dir, "-y"], check=True)
        except FileNotFoundError as e:
            raise RuntimeError("7z (7-Zip) not found on PATH. Please install 7-Zip or add it to PATH.") from e
        # Cleanup parts
        os.remove(zip1)
        os.remove(zip2)
        # Sanity check
        csv_cnt = len(glob(os.path.join(dataset_dir, "*.csv")))
        assert csv_cnt == 3072, f"Expected 3072 csv files, found {csv_cnt}"
        print(f"Verified the extraction of {csv_cnt} csv files.")
    else:
        print(f"Using already installed dataset {dataset_dir} ...")
    return dataset_dir


def _load_single_csv(path: str) -> Tuple[pd.DataFrame, dict]:
    filename = os.path.basename(path)
    match = re.match(_FILENAME_PATTERN, filename)
    assert match, f"Did not find a match for file {filename}"
    meta = match.groupdict()
    df = pd.read_csv(path)
    # remove extra columns
    df = df.loc[:, ~df.columns.str.contains('^Unnamed')]
    cols_to_drop = [
        "ParticipantID", "study_session", "repetition", "timestamp_ms", "phase",
        "HeightNormalization", "ArmLengthNormalization"
    ]
    cols_to_drop += [col for col in df.columns if
                     col.lower().startswith(("ball", "arrow", "rightcontrolleranchor",
                                             "leftcontrolleranchor", "scenario"))]
    df.drop(columns=cols_to_drop, inplace=True, errors='raise')
    return df, meta


def _preprocess_sample(df: pd.DataFrame) -> pd.DataFrame:
    from sklearn.preprocessing import MinMaxScaler
    # local vectors head -> virt. hands
    for lr in ["Left", "Right"]:
        for XYZ in ["X", "Y", "Z"]:
            df[f"{lr}VirtualHand_pos_{XYZ}"] -= df[f"CenterEyeAnchor_pos_{XYZ}"]
    df.drop(columns=[
        'CenterEyeAnchor_pos_X', 'CenterEyeAnchor_pos_Y', 'CenterEyeAnchor_pos_Z',
        'CenterEyeAnchor_euler_X', 'CenterEyeAnchor_euler_Y', 'CenterEyeAnchor_euler_Z'
    ], inplace=True)
    # make everything relative to the beginning
    for c in df.columns:
        df[c] -= df[c].iloc[0]
    # Min Max Scaling
    cols = df.columns
    arr = MinMaxScaler().fit_transform(df)
    return pd.DataFrame(arr, columns=cols)


def get_chi2021_windowsliced_dataset(
    window_size: int = 72,
    stride: int = 72//2,
    n_jobs: int = -1,
    dataset_path: Optional[str] = None,
    reinstall: bool = False,
    label_key: str = 'pid',
    shuffle_slices: bool = False,
) -> WindowslicedDataset:
    """
    Download (if necessary), preprocess and return the PyCHI 2021 dataset
    as a useckit WindowslicedDataset.

    - Samples are labeled by `label_key` parsed from filename (default: 'pid').
    - Only a single set (trainset) is constructed; validation/test can be split
      by the caller if needed.
    """
    # Resolve dataset directory
    base_dir = os.path.dirname(__file__)
    dataset_dir = dataset_path or os.path.join(base_dir, "resources", _DEF_DIRNAME)
    dataset_dir = _download_dataset(dataset_dir, reinstall=reinstall)

    file_paths = glob(os.path.join(dataset_dir, "*.csv"))

    def process_path(p: str):
        df, meta = _load_single_csv(p)
        df = _preprocess_sample(df)
        meta['pid'] = f"P{meta['pid']}"
        meta['session'] = int(meta['session'])
        return df, meta

    results = Parallel(n_jobs=n_jobs)(delayed(process_path)(p) for p in file_paths)

    x_train = [r[0].to_numpy() for r in results if r[1]['session'] == 1 and r[1]['condition'] == "WithoutNormalization"]
    y_train = [r[1]['pid'] for r in results if r[1]['session'] == 1 and r[1]['condition'] == "WithoutNormalization"]
    x_test = [r[0].to_numpy() for r in results if r[1]['session'] == 2 and r[1]['condition'] == "WithoutNormalization"]
    y_test = [r[1]['pid'] for r in results if r[1]['session'] == 2 and r[1]['condition'] == "WithoutNormalization"]

    ds = WindowslicedDataset(
        window_slicing_stride=stride,
        window_slicing_size=window_size,
        trainset_data=x_train,
        trainset_labels=y_train,
        testset_enrollment_data=x_train,
        testset_enrollment_labels=y_train,
        testset_matching_data=x_test,
        testset_matching_labels=y_test,
        normalization_check=False,
        shuffle_window_slices=shuffle_slices,
    )
    return ds


def get_chi2021_dataset(
    n_jobs: int = -1,
    dataset_path: Optional[str] = None,
    reinstall: bool = False,
    label_key: str = 'pid',
    filter: str = 'WithoutNormalization',
    padding: bool = False,
    feature_vector: bool = False,
    ws_enable: bool = False,
    ws_stride: int = 36,
    ws_size: int = 72,
) -> Dataset:
    """
    Download (if necessary), preprocess and return the CHI 2021 dataset
    as a regular Dataset (no window slicing), without any padding.

    - Samples are labeled by `label_key` parsed from filename (default: 'pid').
    - Train/test split follows the same rule as the windowsliced variant:
      session==1 -> enrollment (train), session==2 -> matching (test), both with condition=="WithoutNormalization".
    - Returns variable-length sequences as lists of np.ndarray (float32), one array per sample.
    """
    # Resolve dataset directory
    base_dir = os.path.dirname(__file__)
    dataset_dir = dataset_path or os.path.join(base_dir, "resources", _DEF_DIRNAME)
    dataset_dir = _download_dataset(dataset_dir, reinstall=reinstall)

    file_paths = glob(os.path.join(dataset_dir, "*.csv"))

    def process_path(p: str):
        df, meta = _load_single_csv(p)
        df = _preprocess_sample(df)
        meta['pid'] = f"P{meta['pid']}"
        meta['session'] = int(meta['session'])
        return df, meta

    results = Parallel(n_jobs=n_jobs)(delayed(process_path)(p) for p in file_paths)

    # Filter by condition and session
    def _sel(res, session):
        return [(r[0], r[1][label_key]) for r in res if r[1]['session'] == session and r[1]['condition'] == filter]

    train_pairs = _sel(results, session=1)
    test_pairs = _sel(results, session=2)

    if not train_pairs and not test_pairs:
        raise RuntimeError("No samples found for CHI2021 dataset with the given filters.")

    # Build lists of arrays (no padding)
    x_train = [df for df, _ in train_pairs]
    y_train = np.array([lbl for _, lbl in train_pairs], dtype=str)
    x_test = [df for df, _ in test_pairs]
    y_test = np.array([lbl for _, lbl in test_pairs], dtype=str)

    def q3(data):
        return float(np.percentile(data, 75))

    if padding:
        q3_len: int = int(q3([len(l) for l in x_train + x_test]))
        x_train = pad_sequences(x_train, maxlen=q3_len, dtype=np.float32, padding='pre')
        x_test = pad_sequences(x_test, maxlen=q3_len, dtype=np.float32, padding='pre')

    def calculate_featur_vector(df: pd.DataFrame):
        vector = []
        for c in df.columns:
            vector.append(df[c].min())
            vector.append(df[c].max())
            vector.append(df[c].mean())
            vector.append(df[c].std())
        return vector

    if feature_vector:
        x_train = [calculate_featur_vector(x) for x in x_train]
        x_test = [calculate_featur_vector(x) for x in x_test]

    if ws_enable:
        ds = WindowslicedDataset(
            window_slicing_stride=ws_stride,
            window_slicing_size=ws_size,
            trainset_data=x_train,
            trainset_labels=y_train,
            testset_enrollment_data=x_train,
            testset_enrollment_labels=y_train,
            testset_matching_data=x_test,
            testset_matching_labels=y_test,
            normalization_check=False,)
    else:
        ds = Dataset(
            trainset_data=x_train,
            trainset_labels=y_train,
            testset_enrollment_data=x_train,
            testset_enrollment_labels=y_train,
            testset_matching_data=x_test,
            testset_matching_labels=y_test,
            normalization_check=False,
        )

    return ds
