"""
Reads the files of the HUST Motor dataset and parses them into a usable format for machine learning tasks.
Uses the tag and frequency information to create a structured dataset.
"""

from enum import Enum
from os import listdir, makedirs
from os.path import exists, getmtime, join
import pickle

from dataset.datasetParser import DatasetParser

class MUSTMotorProblemType(Enum):
    HEALTHY = "H"
    BEARING_FAULT = "BF"
    BOWED_ROTOR = "BOW"
    BROKEN_ROTOR_BARS = "BROKEN"
    ROTOR_MISALIGNMENT = "MISAL"
    VOLTAGE_UNBALANCE = "UNBAL"

class MUSTMotorFrequency(Enum):
    FREQ_5HZ = "5HZ"
    FREQ_10HZ = "10HZ"
    FREQ_20HZ = "20HZ"
    FREQ_30HZ = "30HZ"
    

class MUSTMotorDatasetParser(DatasetParser):
    DATASET_PATH = "dataset/HUSTMotor"
    CACHE_PATH = "dataset/.cache/hust_motor_dataset.pkl"
    DATA_COLUMNS = ("time", "x", "y", "z", "sound")
    
    def __init__(self):
        super().__init__()

    def get_files(self):
        """
        Retrieves the list of dataset files from the specified dataset path.
        """
        files = listdir(self.DATASET_PATH)
        # Get the text files
        text_files = sorted(f for f in files if f.endswith('.txt'))
        return text_files
    
    def get_label_from_filename(self, filename):
        """
        Extracts the label (problem type) from the given filename.
        """
        filename = filename.split('.')[0]  # Remove the file extension
        filename = filename.split('_')
        problem_type = MUSTMotorProblemType(filename[0]).value
        frequency = MUSTMotorFrequency(filename[1]).value
        return problem_type, frequency

    def parse_dataset(self):
        """
        Parses the dataset files into one record per measurement row.

        Each record contains the time, three vibration axes, sound, and the
        problem type and frequency encoded in the source filename.
        """
        files = self.get_files()
        if self._cache_is_current(files):
            with open(self.CACHE_PATH, "rb") as cache_file:
                print(f"Loading dataset from cache: {self.CACHE_PATH}")
                return pickle.load(cache_file)

        dataset = []
        for file in files:
            problem_type, frequency = self.get_label_from_filename(file)
            data_started = False
            with open(join(self.DATASET_PATH, file), encoding="utf-8") as dataset_file:
                for line_number, line in enumerate(dataset_file, start=1):
                    if line.strip() == "Time (seconds) and Data Channels":
                        data_started = True
                        continue

                    if not data_started or not line.strip():
                        continue

                    values = line.split()
                    if len(values) != len(self.DATA_COLUMNS):
                        raise ValueError(
                            f"Expected {len(self.DATA_COLUMNS)} values in {file} "
                            f"at line {line_number}, got {len(values)}"
                        )

                    record = dict(zip(self.DATA_COLUMNS, map(float, values)))
                    record["problem_type"] = problem_type
                    record["frequency"] = frequency
                    dataset.append(record)

        makedirs(join(self.DATASET_PATH, "..", ".cache"), exist_ok=True)
        with open(self.CACHE_PATH, "wb") as cache_file:
            pickle.dump(dataset, cache_file, protocol=pickle.HIGHEST_PROTOCOL)

        return dataset

    def get_dataset_motor_problems(self):
        """
        Retrieves the list of unique problem types present in the dataset.
        """
        files = self.get_files()
        problems = set()
        for file in files:
            problem_type, _ = self.get_label_from_filename(file)
            problems.add(problem_type)
        return sorted(problems)

    def _cache_is_current(self, files):
        if not exists(self.CACHE_PATH):
            return False

        cache_timestamp = getmtime(self.CACHE_PATH)
        return all(
            cache_timestamp >= getmtime(join(self.DATASET_PATH, file))
            for file in files
        )
            
    def get_raw_sampling_rate(self):
        return 25600.0  # 25.6 kHz from dataset documentation