import json

import numpy as np
import torch

from transformers import PreTrainedTokenizer
from transformers.tokenization_utils_base import BatchEncoding


class YieldTokenizer(PreTrainedTokenizer):
    """
    Adapter tokenizer for FlexServe's built-in text-classification pipeline.

    This does NOT tokenize natural language.

    It accepts:
      - a JSON string, or
      - the yield input dictionary

    and converts it into:
      weather:     [B, 52, 6]
      soil:        [B, 66]
      crop_id:     [B]
      horizon_idx: [B]

    Normalization is intentionally NOT performed here.
    The sequence-classification model wrapper performs normalization
    using statistics stored in config.json.
    """

    vocab_files_names = {}
    model_input_names = [
        "weather",
        "soil",
        "crop_id",
        "horizon_idx",
    ]

    def __init__(
        self,
        weather_vars=None,
        soil_vars=None,
        K=52,
        eval_cutoffs=None,
        **kwargs,
    ):
        self.weather_vars = list(weather_vars or [])
        self.soil_vars = list(soil_vars or [])
        self.K = int(K)
        self.eval_cutoffs = list(
            eval_cutoffs
            or [20, 24, 28, 32, 36, 40, 44, 48, 52]
        )

        super().__init__(
            pad_token="[PAD]",
            unk_token="[UNK]",
            **kwargs,
        )

    @property
    def vocab_size(self):
        return 2

    def get_vocab(self):
        return {
            "[PAD]": 0,
            "[UNK]": 1,
        }

    def _tokenize(self, text, **kwargs):
        return ["[UNK]"]

    def _convert_token_to_id(self, token):
        return 0 if token == "[PAD]" else 1

    def _convert_id_to_token(self, index):
        return "[PAD]" if index == 0 else "[UNK]"

    def save_vocabulary(self, save_directory, filename_prefix=None):
        return ()

    def _parse_sample(self, sample):

        if isinstance(sample, str):
            sample = sample.strip()

            try:
                sample = json.loads(sample)
            except json.JSONDecodeError as exc:
                raise ValueError(
                    "Input must be a valid JSON string."
                ) from exc

            # Handle a JSON string containing another JSON string.
            if isinstance(sample, str):
                try:
                    sample = json.loads(sample)
                except json.JSONDecodeError as exc:
                    raise ValueError(
                        "Input string does not contain valid yield JSON."
                    ) from exc

        if not isinstance(sample, dict):
            raise ValueError(
                "Yield input must be a JSON object/dictionary."
            )

        crop = str(
            sample.get("crop", "corn")
        ).strip().lower()

        if crop not in ("corn", "maize"):
            raise ValueError(
                "This released model supports corn only."
            )

        if "weather" not in sample:
            raise ValueError(
                "Missing 'weather' object."
            )

        if "soil" not in sample:
            raise ValueError(
                "Missing 'soil' object."
            )

        weather_dict = sample["weather"]
        soil_dict = sample["soil"]

        # --------------------------------------------
        # Weather: [52, 6]
        # --------------------------------------------

        weather_cols = []

        for var in self.weather_vars:

            if var not in weather_dict:
                raise ValueError(
                    f"Missing weather variable '{var}'."
                )

            values = np.asarray(
                weather_dict[var],
                dtype=np.float32,
            )

            if values.ndim != 1:
                raise ValueError(
                    f"Weather '{var}' must be one-dimensional."
                )

            if len(values) != self.K:
                raise ValueError(
                    f"Weather '{var}' requires exactly "
                    f"{self.K} weekly values; received {len(values)}."
                )

            weather_cols.append(values)

        weather = np.stack(
            weather_cols,
            axis=1,
        ).astype(np.float32)

        # --------------------------------------------
        # Soil: [66]
        # --------------------------------------------

        soil = []

        for var in self.soil_vars:

            if var not in soil_dict:
                raise ValueError(
                    f"Missing soil variable '{var}'."
                )

            soil.append(
                float(soil_dict[var])
            )

        soil = np.asarray(
            soil,
            dtype=np.float32,
        )

        # --------------------------------------------
        # Cutoff
        # --------------------------------------------

        cutoff = int(
            sample.get(
                "cutoff",
                max(self.eval_cutoffs),
            )
        )

        if cutoff not in self.eval_cutoffs:
            raise ValueError(
                f"Unsupported cutoff {cutoff}. "
                f"Supported cutoffs are {self.eval_cutoffs}."
            )

        return {
            "weather": weather,
            "soil": soil,
            "crop_id": 0,
            "horizon_idx": cutoff,
        }

    def __call__(
        self,
        text=None,
        text_pair=None,
        return_tensors=None,
        **kwargs,
    ):
        # --------------------------------------------------
        # FlexServe/HF may call tokenizer(**input_dict)
        # instead of tokenizer(json_string).
        # --------------------------------------------------

        if text is None and "weather" in kwargs and "soil" in kwargs:

            sample = {
                "crop": kwargs.pop("crop", "corn"),
                "weather": kwargs.pop("weather"),
                "soil": kwargs.pop("soil"),
                "cutoff": kwargs.pop(
                    "cutoff",
                    max(self.eval_cutoffs),
                ),
            }

            samples = [sample]

        elif isinstance(text, (list, tuple)):

            samples = list(text)

        else:

            samples = [text]

        parsed = [
            self._parse_sample(sample)
            for sample in samples
        ]

        weather = torch.tensor(
            np.stack(
                [x["weather"] for x in parsed],
                axis=0,
            ),
            dtype=torch.float32,
        )

        soil = torch.tensor(
            np.stack(
                [x["soil"] for x in parsed],
                axis=0,
            ),
            dtype=torch.float32,
        )

        crop_id = torch.tensor(
            [
                x["crop_id"]
                for x in parsed
            ],
            dtype=torch.long,
        )

        horizon_idx = torch.tensor(
            [
                x["horizon_idx"]
                for x in parsed
            ],
            dtype=torch.long,
        )

        return BatchEncoding(
            {
                "weather": weather,
                "soil": soil,
                "crop_id": crop_id,
                "horizon_idx": horizon_idx,
            }
        )