import os
from pathlib import Path
from typing import Any, BinaryIO, Mapping, Optional, Union
import torch

from config import default_config
from featex import load_audio, Preprocessor
from model import Classifier

class Pipeline:
    def __init__(self, checkpoint: Optional[str | Path] = None, config: Optional[Mapping[str, Any]] = None, device: Optional[torch.device] = None):
        if checkpoint is None:
            file_dir = Path(__file__).parent.resolve()
            checkpoint = file_dir / "dam3.1.ckpt"
        if config is None:
            config = default_config
        if device is None:
            if torch.cuda.is_available():
                device = torch.device("cuda:0")
            else:
                device = torch.device("cpu")
        self.device = device
        self.model = Classifier(**config)
        self.preprocessor = Preprocessor(**self.model.preprocessor_config)
        state_dict = torch.load(checkpoint, map_location=device)
        self.model.load_state_dict(state_dict)
        self.model.to(self.device)
        self.model.eval()

    def run_on_features(self, features: torch.Tensor, quantize: bool = True):
        scores = self.model(features, torch.tensor([features.shape[0]], device=self.device))[0]
        if quantize:
            return {k: int(v.item()) for k, v in self.model.quantize_scores(scores).items()}
        else:
            return scores

    def run_on_audio(self, audio: torch.Tensor, quantize: bool = True):
        features = self.preprocessor.preprocess_with_audio_normalization(audio)
        return self.run_on_features(features.to(self.device), quantize=quantize)

    def run_on_file(self, source: Union[BinaryIO, str, os.PathLike], quantize=True):
        audio = load_audio(source)
        return self.run_on_audio(audio, quantize=quantize)