import ast
import re
from collections.abc import Iterable, Mapping

import torch
import torch.nn as nn
from huggingface_hub import PyTorchModelHubMixin


class DataProcessor:

    def __init__(self):
        self.vocab = None

    def _check_readiness(self) -> bool:
        if self.vocab is None:
            raise ValueError('Missing vocab. Ensure loading respective vocab.json.')

    @staticmethod
    def _parse_tags(tags) -> list:
        """Read tag lists, including list strings loaded from CSV."""
        if isinstance(tags, Mapping):
            return []
        if isinstance(tags, str):
            control_escapes = {code: f'\\x{code:02x}' for code in range(32)}
            values = []
            for match in re.finditer('\'(?:\\\\.|[^\'\\\\])*\'|"(?:\\\\.|[^"\\\\])*"', tags):
                try:
                    values.append(ast.literal_eval(match.group(0).translate(control_escapes)))
                except (SyntaxError, ValueError):
                    continue
        elif isinstance(tags, Iterable):
            values = tags
        else:
            return []
        return values

    @staticmethod
    def preprocess_tags(tags) -> list[str]:
        """Normalize tags using the vocabulary extraction rules."""
        cleaned = []
        for value in DataProcessor._parse_tags(tags):
            if not isinstance(value, str):
                continue
            tag = value.lower()
            for old, new in ((' - ', ''), ('-', ''), ('.', ''), ('|', ''), ('ä', 'a'), ('ö', 'o'), ('ü', 'u')):
                tag = tag.replace(old, new)
            if tag:
                cleaned.append(tag)
        return cleaned

    def normalize_tags(self, tags) -> list[str]:
        """Preprocess raw or normalized tags and keep only known tags."""
        self._check_readiness()
        return [tag for tag in self.preprocess_tags(tags) if tag in self.known_tags_dict]

    def vectorize_tags(self, tags) -> torch.Tensor:
        """Creates onehot-encoded vector of used tags for classification.
        Ignores unknown tags.

        Returns vector full of zeros if no valid tags were found or tags to normalize were [] or `None`"""
        self._check_readiness()
        base = torch.zeros(self.nr_of_known_tags, dtype=torch.float32)
        for p_tag in self.normalize_tags(tags):
            base[int(self.known_tags_dict[p_tag])] = 1
        return base

    def process_api(self, response_dict: dict) -> torch.Tensor:
        """Vectorize an OFF product dictionary; missing tags produce a zero vector."""
        self._check_readiness()
        return self.vectorize_tags(response_dict.get('categories_tags'))


class MLP(nn.Module, PyTorchModelHubMixin):
    def __init__(self, input_features:int):
        super().__init__()

        self.class_mapping = {
            0: "Obst & Gemüse",
            1: "Fleisch, Fisch & Alternativen",
            2: "Milchprodukte & Eier",
            3: "Brot & Backwaren",
            4: "Nudeln, Reis & Getreide",
            5: "Fertigprodukte",
            6: "Backzutaten",
            7: "Getränke",
            8: "Snacks & Süßigkeiten",
            9: "Öle, Gewürze & Saucen",
            10: "Brotaufstriche",
            11: "Sonstiges"
        }

        self.input_layer = nn.Linear(in_features=input_features, out_features=2000)     # classes 0-11
        self.hidden1 = nn.Linear(self.input_layer.out_features, 1000)
        self.hidden2 = nn.Linear(self.hidden1.out_features, 100)
        self.hidden3 = nn.Linear(self.hidden2.out_features, 12)

        self.softmax = nn.Softmax(1)
        self.leaky_r = nn.LeakyReLU(inplace=True)
        self.dropout1 = nn.Dropout(p=.2)
        self.dropout2 = nn.Dropout(p=.2)

        self.forward_seq = nn.Sequential(
            self.input_layer,
            self.leaky_r,

            self.hidden1,
            self.dropout1,
            self.leaky_r,

            self.hidden2,
            self.dropout2,
            self.leaky_r,

            self.hidden3
        )

    def forward(self, X:torch.Tensor) -> torch.Tensor:
        """Performs forward pass."""
        return self.forward_seq(X)

    def predict(self, X, string_classes:bool=True) -> torch.Tensor:
        if X.ndim == 1:
            X = X.unsqueeze(0)
        preds = self.forward_seq(X)
        res = int(torch.argmax(preds, dim=1).cpu())
        return self.class_mapping[res] if string_classes else res


class FoodClassifier(MLP):
    """MLP with preprocessing and the standard Hugging Face Hub mixin."""

    def __init__(self, vocab: list[str]):
        super().__init__(input_features=len(vocab))
        # The mixin saves the ordered vocabulary in config.json.
        self.processor = DataProcessor()
        self.processor.vocab = list(vocab)
        self.processor.known_tags_dict = {tag: idx for idx, tag in enumerate(vocab)}
        self.processor.nr_of_known_tags = len(vocab)

    @torch.inference_mode()
    def predict(self, product: dict) -> str:
        """Return the category name for a product containing categories_tags."""
        vector = self.processor.process_api(product).to(next(self.parameters()).device)
        return super().predict(vector, string_classes=True)