# Copyright (c) 2025 SciKnowOrg
#
# Licensed under the MIT License (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      https://opensource.org/licenses/MIT
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import os
import re
import random

import pandas as pd
import torch
import Levenshtein
from datasets import Dataset
from typing import Any, Optional, List, Tuple, Dict
from transformers import (
    AutoTokenizer,
    AutoModelForSequenceClassification,
    AutoModelForCausalLM,
    BertTokenizer,
    BertForSequenceClassification,
    pipeline,
    Trainer,
    TrainingArguments,
)
from tqdm import tqdm
from ...base import AutoLearner, AutoPrompt
from ...utils import taxonomy_split, train_test_split as ontology_split
from ...data_structure import OntologyData, TaxonomicRelation


class SKHNLPTaxonomyPrompts(AutoPrompt):
    """Builds the 7 taxonomy prompts used during fine-tuning / inference.

    The class stores a small inventory of prompt templates that verbalize the
    (parent, child) relationship using different phrasings. Each template ends
    with a masked token slot intended for True/False classification.
    """

    def __init__(self) -> None:
        """Initialize prompt templates and the default prompt in the base class."""
        super().__init__(
            prompt_template="{parent} is the superclass of {child}. This statement is [MASK]."
        )
        self.templates: List[str] = [
            "{parent} is the superclass of {child}. This statement is [MASK].",
            "{child} is a subclass of {parent}. This statement is [MASK].",
            "{parent} is the parent class of {child}. This statement is [MASK].",
            "{child} is a child class of {parent}. This statement is [MASK].",
            "{parent} is a supertype of {child}. This statement is [MASK].",
            "{child} is a subtype of {parent}. This statement is [MASK].",
            "{parent} is an ancestor class of {child}. This statement is [MASK].",
        ]

    def format(self, parent: str, child: str, template_idx: int) -> str:
        """Render a prompt for a (parent, child) pair using a specific template.

        Args:
            parent: The parent/superclass label.
            child: The child/subclass label.
            template_idx: Index into the internal `templates` list.

        Returns:
            The fully formatted prompt string.
        """
        return self.templates[template_idx].format(parent=parent, child=child)


class SKHNLPSequentialFTLearner(AutoLearner):
    """
    BERT-based classifier for taxonomy discovery.

    With OntologyData:
      * TRAIN: ontology-aware split; create balanced train/eval with negatives.
      * PREDICT/TEST: notebook-style parent selection -> list[{'parent', 'child'}].

    With DataFrame/list:
      * TRAIN: taxonomy_split + negatives; build prompts and fine-tune.
      * PREDICT/TEST: pairwise binary classification (returns label + score).
    """

    def __init__(
            self,
            # core
            model_name: str = "bert-large-uncased",
            n_prompts: int = 7,
            random_state: int = 1403,
            num_labels: int = 2,
            device: str = "cpu",  # "cuda" | "cpu" | None (auto)
            # data split & negative sampling (now configurable)
            eval_fraction: float = 0.16,
            neg_ratio_reversed: float = 1 / 3,
            neg_ratio_manipulated: float = 2 / 3,
            # ---- expose TrainingArguments as individual user-defined args ----
            output_dir: str = "./results/",
            num_train_epochs: int = 1,
            per_device_train_batch_size: int = 4,
            per_device_eval_batch_size: int = 4,
            warmup_steps: int = 500,
            weight_decay: float = 0.01,
            logging_dir: str = "./logs/",
            logging_steps: int = 50,
            eval_strategy: str = "epoch",
            save_strategy: str = "epoch",
            load_best_model_at_end: bool = True,
            use_fast_tokenizer: Optional[bool] = None,
            trust_remote_code: bool = False,
    ) -> None:
        """Configure the sequential fine-tuning learner.

        Args:
            model_name: HF model id or local path for the BERT backbone.
            n_prompts: Number of prompt variants to iterate over sequentially.
            random_state: RNG seed for shuffling/sampling steps.
            num_labels: Number of classes for the classifier head.
            device: Force device ('cuda' or 'cpu'). If None, auto-detects CUDA.
            eval_fraction: Fraction of positives to hold out for evaluation.
            neg_ratio_reversed: Proportion of reversed-parent negatives vs positives.
            neg_ratio_manipulated: Proportion of random-parent negatives vs positives.
            output_dir: Directory where HF Trainer writes checkpoints/outputs.
            num_train_epochs: Number of epochs per prompt.
            per_device_train_batch_size: Training batch size per device.
            per_device_eval_batch_size: Evaluation batch size per device.
            warmup_steps: Linear warmup steps for LR scheduler.
            weight_decay: Weight decay coefficient.
            logging_dir: Directory for Trainer logs.
            logging_steps: Interval for log events (in steps).
            eval_strategy: Evaluation schedule ('no', 'steps', 'epoch').
            save_strategy: Checkpoint save schedule ('no', 'steps', 'epoch').
            load_best_model_at_end: Whether to restore the best checkpoint.
            use_fast_tokenizer: Force fast/slow tokenizer. If None, try fast then fallback to slow.
        Notes:
            The model is fine-tuned *sequentially* across prompt columns.
            You can control the eval split and negative sampling mix via
            `eval_fraction`, `neg_ratio_reversed`, and `neg_ratio_manipulated`.
        """
        super().__init__()
        self.model_name = model_name
        self.n_prompts = n_prompts
        self.random_state = random_state
        self.num_labels = num_labels
        self.device = device

        # user-tunable ratios / split
        self._eval_fraction = float(eval_fraction)
        self._neg_ratio_reversed = float(neg_ratio_reversed)
        self._neg_ratio_manipulated = float(neg_ratio_manipulated)
        if not (0.0 < self._eval_fraction < 1.0):
            raise ValueError("eval_fraction must be in (0, 1).")
        if self._neg_ratio_reversed < 0 or self._neg_ratio_manipulated < 0:
            raise ValueError("neg_ratio_* must be >= 0.")

        self.tokenizer: Optional[BertTokenizer] = None
        self.model: Optional[BertForSequenceClassification] = None
        self.prompter = SKHNLPTaxonomyPrompts()

        # Candidate parents (unique parent list) for multi-class parent selection.
        self._candidate_parents: Optional[List[str]] = None

        # Keep last train/eval tables for inspection
        self._last_train: Optional[pd.DataFrame] = None
        self._last_eval: Optional[pd.DataFrame] = None
        self.trust_remote_code = bool(trust_remote_code)
        self.use_fast_tokenizer = use_fast_tokenizer
        self.per_device_eval_batch_size = per_device_eval_batch_size

        random.seed(self.random_state)

        # Build TrainingArguments from the individual user-defined values
        self.training_args = TrainingArguments(
            output_dir=output_dir,
            num_train_epochs=num_train_epochs,
            per_device_train_batch_size=per_device_train_batch_size,
            per_device_eval_batch_size=per_device_eval_batch_size,
            warmup_steps=warmup_steps,
            weight_decay=weight_decay,
            logging_dir=logging_dir,
            logging_steps=logging_steps,
            eval_strategy=eval_strategy,
            save_strategy=save_strategy,
            load_best_model_at_end=load_best_model_at_end,
        )

    def load(self, model_id: Optional[str] = None, **_: Any) -> None:
        """Load tokenizer & model in a backbone-agnostic way; move model to self.device."""
        model_id = model_id or self.model_name

        # ---- Tokenizer (robust fast→slow fallback unless explicitly set) ----
        if self.use_fast_tokenizer is None:
            try:
                self.tokenizer = AutoTokenizer.from_pretrained(
                    model_id, use_fast=True, trust_remote_code=self.trust_remote_code
                )
            except Exception as fast_err:
                print(
                    f"[tokenizer] Fast tokenizer failed: {fast_err}. Falling back to slow tokenizer..."
                )
                self.tokenizer = AutoTokenizer.from_pretrained(
                    model_id, use_fast=False, trust_remote_code=self.trust_remote_code
                )
        else:
            self.tokenizer = AutoTokenizer.from_pretrained(
                model_id,
                use_fast=self.use_fast_tokenizer,
                trust_remote_code=self.trust_remote_code,
            )

        # Ensure pad token exists (some models lack it)
        if getattr(self.tokenizer, "pad_token", None) is None:
            # Try sensible fallbacks
            fallback = (
                    getattr(self.tokenizer, "eos_token", None)
                    or getattr(self.tokenizer, "sep_token", None)
                    or getattr(self.tokenizer, "cls_token", None)
            )
            if fallback is not None:
                self.tokenizer.pad_token = fallback

        # ---- Model (classifier head sized to self.num_labels) ----
        self.model = AutoModelForSequenceClassification.from_pretrained(
            model_id,
            num_labels=self.num_labels,
            trust_remote_code=self.trust_remote_code,
            # Allows swapping in a new head size even if the checkpoint differs
            ignore_mismatched_sizes=True,
        )

        # Make sure padding ids line up
        if (
                getattr(self.model.config, "pad_token_id", None) is None
                and getattr(self.tokenizer, "pad_token_id", None) is not None
        ):
            self.model.config.pad_token_id = self.tokenizer.pad_token_id

        # Set problem type (single-label classification by default)
        # If you plan multi-label, you'd switch to "multi_label_classification"
        self.model.config.problem_type = "single_label_classification"

        # Move to target device
        self.model.to(self.device)

    def tasks_ground_truth_former(self, data: Any, task: str) -> Any:
        """Normalize ground-truth inputs for 'taxonomy-discovery'.

        Supports DataFrame with columns ['parent','child',('label')],
        list of dicts, or falls back to the base class behavior.

        Args:
            data: Input object to normalize.
            task: Task name, passed from the outer pipeline.

        Returns:
            A list of dictionaries with keys 'parent', 'child', and optionally
            'label' when present in the input.
        """
        if task != "taxonomy-discovery":
            return super().tasks_ground_truth_former(data, task)

        if isinstance(data, pd.DataFrame):
            if "label" in data.columns:
                return [
                    {"parent": p, "child": c, "label": bool(lbl)}
                    for p, c, lbl in zip(data["parent"], data["child"], data["label"])
                ]
            return [
                {"parent": p, "child": c} for p, c in zip(data["parent"], data["child"])
            ]

        if isinstance(data, list):
            return data

        return super().tasks_ground_truth_former(data, task)

    def _make_negatives(
            self, positives_df: pd.DataFrame
    ) -> Tuple[pd.DataFrame, pd.DataFrame]:
        """Create two types of negatives from a positives table.

        Returns:
            A tuple `(reversed_df, manipulated_df)` where:
                - `reversed_df`: pairs with parent/child columns swapped, label=False.
                - `manipulated_df`: pairs with the parent replaced by a random
                  *different* parent from the same pool, label=False.

        Notes:
            The input DataFrame must contain columns ['parent', 'child'].
        """
        unique_parents = positives_df["parent"].unique().tolist()

        def as_reversed(df: pd.DataFrame) -> pd.DataFrame:
            out = df.copy()
            out[["parent", "child"]] = out[["child", "parent"]].values
            out["label"] = False
            return out

        def with_random_parent(df: pd.DataFrame) -> pd.DataFrame:
            def pick_other_parent(p: str) -> str:
                pool = [x for x in unique_parents if x != p]
                return random.choice(pool) if pool else p

            out = df.copy()
            out["parent"] = out["parent"].apply(pick_other_parent)
            out["label"] = False
            return out

        return as_reversed(positives_df), with_random_parent(positives_df)

    def _balance_with_negatives(
            self,
            positives_df: pd.DataFrame,
            reversed_df: pd.DataFrame,
            manipulated_df: pd.DataFrame,
    ) -> pd.DataFrame:
        """Combine positives with negatives using configured ratios.

        Sampling ratios are defined by the instance settings
        `self._neg_ratio_reversed` and `self._neg_ratio_manipulated`,
        keeping the positives count unchanged.

        Args:
            positives_df: Positive pairs with `label=True`.
            reversed_df: Negative pairs produced by flipping parent/child.
            manipulated_df: Negative pairs with randomly reassigned parents.

        Returns:
            A deduplicated, shuffled DataFrame with a class-balanced mix.
        """
        n_pos = len(positives_df)
        n_rev = int(n_pos * self._neg_ratio_reversed)
        n_man = int(n_pos * self._neg_ratio_manipulated)

        combined = pd.concat(
            [
                positives_df.sample(n_pos, random_state=self.random_state),
                reversed_df.sample(n_rev, random_state=self.random_state),
                manipulated_df.sample(n_man, random_state=self.random_state),
            ],
            ignore_index=True,
        )
        combined = combined.drop_duplicates(
            subset=["parent", "child", "label"]
        ).reset_index(drop=True)
        return combined

    def _add_prompt_columns(self, df: pd.DataFrame) -> pd.DataFrame:
        """Append one column per prompt variant to the given pairs table.

        For each row `(parent, child)`, creates columns `prompt_1 ... prompt_n`.

        Args:
            df: Input DataFrame with columns ['parent', 'child', ...].

        Returns:
            A copy of `df` including the newly added prompt columns.
        """
        out = df.copy()
        for i in range(self.n_prompts):
            out[f"prompt_{i + 1}"] = out.apply(
                lambda r, k=i: self.prompter.format(r["parent"], r["child"], k), axis=1
            )
        return out

    def _df_from_relations(
            self, relations: List[TaxonomicRelation], label: bool = True
    ) -> pd.DataFrame:
        """Convert a list of `TaxonomicRelation` to a DataFrame.

        Args:
            relations: Iterable of `TaxonomicRelation(parent, child)`.
            label: Class label to assign to all resulting rows.

        Returns:
            DataFrame with columns ['parent', 'child', 'label'].
        """
        if not relations:
            return pd.DataFrame(columns=["parent", "child", "label"])
        return pd.DataFrame(
            [{"parent": r.parent, "child": r.child, "label": label} for r in relations]
        )

    def _relations_from_df(self, df: pd.DataFrame) -> List[TaxonomicRelation]:
        """Convert a DataFrame to a list of `TaxonomicRelation`.

        Args:
            df: DataFrame with columns ['parent', 'child'].

        Returns:
            List of `TaxonomicRelation` objects in row order.
        """
        return [
            TaxonomicRelation(parent=p, child=c)
            for p, c in zip(df["parent"], df["child"])
        ]

    def _build_masked_prompt(
            self, parent: str, child: str, index_1_based: int, mask_token: str = "[MASK]"
    ) -> str:
        """Construct one of several True/False prompts with a mask token.

        Args:
            parent: Parent label.
            child: Child label.
            index_1_based: 1-based index selecting a template.
            mask_token: The token used to denote the masked label.

        Returns:
            A formatted prompt string.
        """
        prompts_1based = [
            f"{parent} is the superclass of {child}. This statement is {mask_token}.",
            f"{child} is a subclass of {parent}. This statement is {mask_token}.",
            f"{parent} is the parent class of {child}. This statement is {mask_token}.",
            f"{child} is a child class of {parent}. This statement is {mask_token}.",
            f"{parent} is a supertype of {child}. This statement is {mask_token}.",
            f"{child} is a subtype of {parent}. This statement is {mask_token}.",
            f"{parent} is an ancestor class of {child}. This statement is {mask_token}.",
            f"{child} is a descendant classs of {child}. This statement is {mask_token}.",
            f'"{parent}" is the superclass of "{child}". This statement is {mask_token}.',
        ]
        return prompts_1based[index_1_based - 1]

    @torch.no_grad()
    def _predict_prompt_true_false(self, sentence: str) -> bool:
        """Run a single True/False prediction on a prompt.

        Args:
            sentence: Fully formatted prompt text.

        Returns:
            True iff the predicted class index is 1 (positive).
        """
        enc = self.tokenizer(sentence, return_tensors="pt").to(self.model.device)
        logits = self.model(**enc).logits
        predicted_label = torch.argmax(logits, dim=1).item()
        return predicted_label == 1

    def _select_parent_via_prompts_main(self, child: str) -> str:
        """Select the most likely parent for a given child via prompt voting.

        The procedure:
          1) Generate prompts for each candidate parent at increasing "levels".
          2) Accumulate votes from the True/False classifier.
          3) Resolve ties by recursing to the next level; after 4 levels, break ties randomly.

        Args:
            child: The child label whose parent should be predicted.

        Returns:
            The chosen parent string.

        Raises:
            AssertionError: If candidate parents were not initialized.
        """
        assert self._candidate_parents, "Candidate parents not initialized."
        scores: dict[str, int] = {p: 0 for p in self._candidate_parents}

        def prompt_indices_for_level(level: int) -> List[int]:
            if level == 0:
                return [1]
            return [2 * level, 2 * level + 1]

        def recurse(active_parents: List[str], level: int) -> str:
            idxs = [
                i for i in prompt_indices_for_level(level) if 1 <= i <= self.n_prompts
            ]
            if idxs:
                for parent in active_parents:
                    votes = sum(
                        1
                        for idx in idxs
                        if self._predict_prompt_true_false(
                            self._build_masked_prompt(
                                parent=parent, child=child, index_1_based=idx
                            )
                        )
                    )
                    scores[parent] += votes

            max_score = max(scores[p] for p in active_parents)
            tied = [p for p in active_parents if scores[p] == max_score]
            if len(tied) == 1:
                return tied[0]
            if level < 4:
                return recurse(tied, level + 1)
            return random.choice(tied)

        return recurse(list(scores.keys()), level=0)

    def _select_parent_via_prompts(self, child: str) -> str:
        """Select the most likely parent for a given child via batched inference.

        This vectorized approach processes all candidate parents simultaneously,
        making it O(n_prompts) instead of O(candidates × n_prompts).

        Args:
            child: The child label whose parent should be predicted.

        Returns:
            The chosen parent string with highest average probability.

        Raises:
            AssertionError: If candidate parents were not initialized.
        """
        assert self._candidate_parents, "Candidate parents not initialized."

        # Build all prompts for all candidates at once
        # Shape: (n_candidates, n_prompts)
        all_prompts = []
        for parent in self._candidate_parents:
            parent_prompts = [
                self._build_masked_prompt(parent, child, idx + 1)
                for idx in range(self.n_prompts)
            ]
            all_prompts.extend(parent_prompts)

        # Tokenize all prompts in one batch
        encodings = self.tokenizer(
            all_prompts,
            return_tensors="pt",
            padding=True,
            truncation=True,
        ).to(self.model.device)

        # Single forward pass for all prompts
        logits = self.model(**encodings).logits

        # Get probabilities for the "True" class (index 1)
        true_probs = torch.softmax(logits, dim=1)[:, 1]

        # Reshape to (n_candidates, n_prompts)
        true_probs = true_probs.view(len(self._candidate_parents), self.n_prompts)

        # Average across prompts for each candidate
        avg_scores = true_probs.mean(dim=1)

        # Select candidate with highest average score
        best_idx = torch.argmax(avg_scores).item()

        return self._candidate_parents[best_idx]


    def _select_parents_batch(self,
                              children: List[str],
                              batch_size: int = 8,
                              max_prompts_per_forward: int = 32) -> List[str]:
        """Select parents for multiple children with memory-efficient chunked inference.

        Args:
            children: List of child labels to find parents for.
            batch_size: Number of children to process simultaneously.
            max_prompts_per_forward: Max prompts per forward pass (reduce if OOM: try 32, 16, or 8).

        Returns:
            List of chosen parent strings (same order as input children).
        """
        assert self._candidate_parents, "Candidate parents not initialized."

        all_predictions = []

        for batch_start in tqdm(range(0, len(children), batch_size), desc="Batch parent selection"):
            batch_children = children[batch_start:batch_start + batch_size]

            # Build all prompts for this batch of children
            # Shape: (batch_size × n_candidates × n_prompts)
            batch_prompts = []
            for child in batch_children:
                for parent in self._candidate_parents:
                    for idx in range(self.n_prompts):
                        batch_prompts.append(
                            self._build_masked_prompt(parent, child, idx + 1)
                        )

            # Process prompts in chunks to avoid OOM
            all_true_probs = []
            for chunk_start in range(0, len(batch_prompts), max_prompts_per_forward):
                chunk_prompts = batch_prompts[chunk_start:chunk_start + max_prompts_per_forward]

                # Tokenize chunk
                encodings = self.tokenizer(
                    chunk_prompts,
                    return_tensors="pt",
                    padding=True,
                    max_length=256,
                    truncation=True,
                ).to(self.model.device)

                # Forward pass on chunk
                logits = self.model(**encodings).logits
                true_probs = torch.softmax(logits, dim=1)[:, 1]
                all_true_probs.append(true_probs.cpu())  # Move to CPU to free GPU memory

                # Clear GPU cache
                del encodings, logits
                torch.cuda.empty_cache() if torch.cuda.is_available() else None

            # Concatenate all chunks
            all_true_probs = torch.cat(all_true_probs, dim=0)

            # Reshape: (batch_size, n_candidates, n_prompts)
            all_true_probs = all_true_probs.view(len(batch_children), len(self._candidate_parents), self.n_prompts)
            # Average across prompts, then select best candidate per child
            avg_scores = all_true_probs.mean(dim=2)  # (batch_size, n_candidates)
            best_indices = torch.argmax(avg_scores, dim=1)  # (batch_size,)

            # Map indices to parent names
            batch_predictions = [self._candidate_parents[idx.item()] for idx in best_indices]
            all_predictions.extend(batch_predictions)

        return all_predictions

    def _taxonomy_discovery(self, data: Any, test: bool = False):
        """
        TRAIN:
          - OntologyData -> ontology-aware split; negatives per split; balanced sets.
          - DataFrame/list -> taxonomy_split for positives; negatives proportional.
        TEST:
          - OntologyData -> parent selection: [{'parent': predicted, 'child': child}]
          - DataFrame/list -> binary pair classification with 'label' + 'score'

        Args:
            data: One of {OntologyData, pandas.DataFrame, list[dict], list[tuple]}.
            test: If True, run inference; otherwise perform training.

        Returns:
            - On training: None (model is fine-tuned in-place).
            - On inference with OntologyData: list of {'parent','child'} predictions.
            - On inference with pairs: list of dicts including 'label' and 'score'.
        """
        is_ontology_object = isinstance(data, OntologyData)

        # Normalize input
        if isinstance(data, pd.DataFrame):
            pairs_df = data.copy()
        elif isinstance(data, list):
            pairs_df = pd.DataFrame(data)
        else:
            gt_pairs = super().tasks_ground_truth_former(data, "taxonomy-discovery")
            pairs_df = pd.DataFrame(gt_pairs)
            if "label" not in pairs_df.columns:
                pairs_df["label"] = True

        # Maintain candidate parents across calls
        if "parent" in pairs_df.columns:
            parents_in_call = sorted(pd.unique(pairs_df["parent"]).tolist())
            if test:
                if self._candidate_parents is None:
                    self._candidate_parents = parents_in_call
                else:
                    self._candidate_parents = sorted(
                        set(self._candidate_parents).union(parents_in_call)
                    )
            else:
                if self._candidate_parents is None:
                    self._candidate_parents = parents_in_call

        if test:
            if is_ontology_object and self._candidate_parents:
                # predictions: List[dict[str, str]] = []
                # for _, row in tqdm(pairs_df.iterrows(), desc="Select paren via prompt:"):
                #     child_term = row["child"]
                #     chosen_parent = self._select_parent_via_prompts(child_term)
                #     predictions.append({"parent": chosen_parent, "child": child_term})
                # return predictions
                # predictions: List[dict[str, str]] = []
                children_list = pairs_df["child"].tolist()
                chosen_parents = self._select_parents_batch(children_list, batch_size=self.per_device_eval_batch_size)
                predictions = [{"parent": parent, "child": child}
                               for parent, child in zip(chosen_parents, children_list)]
                return predictions

            # pairwise binary classification
            prompts_df = self._add_prompt_columns(pairs_df.copy())
            true_probs_by_prompt: List[torch.Tensor] = []

            for i in tqdm(range(self.n_prompts), desc="Prompt via prompt:"):
                col = f"prompt_{i + 1}"
                enc = self.tokenizer(
                    prompts_df[col].tolist(),
                    return_tensors="pt",
                    padding=True,
                    truncation=True,
                ).to(self.model.device)
                with torch.no_grad():
                    logits = self.model(**enc).logits
                    true_probs_by_prompt.append(torch.softmax(logits, dim=1)[:, 1])

            avg_true_prob = torch.stack(true_probs_by_prompt, dim=0).mean(0)
            predicted_bool = (avg_true_prob >= 0.5).cpu().tolist()

            results: List[dict[str, Any]] = []
            for p, c, s, yhat in tqdm(zip(pairs_df["parent"],
                                          pairs_df["child"],
                                          avg_true_prob.tolist(),
                                          predicted_bool), desc="Append:"):
                results.append({"parent": p, "child": c, "label": int(bool(yhat)), "score": float(s)})
            return results

        if isinstance(data, OntologyData):
            train_onto, eval_onto = ontology_split(data,
                                                   test_size=self._eval_fraction,
                                                   random_state=self.random_state,
                                                   verbose=False)

            train_pos_rel: List[TaxonomicRelation] = getattr(train_onto.type_taxonomies, "taxonomies", []) or []
            eval_pos_rel: List[TaxonomicRelation] = getattr(eval_onto.type_taxonomies, "taxonomies", []) or []

            train_pos_df = self._df_from_relations(train_pos_rel, label=True)
            eval_pos_df = self._df_from_relations(eval_pos_rel, label=True)

            tr_rev_df, tr_man_df = self._make_negatives(train_pos_df)
            ev_rev_df, ev_man_df = self._make_negatives(eval_pos_df)

            train_df = self._balance_with_negatives(train_pos_df, tr_rev_df, tr_man_df)
            eval_df = self._balance_with_negatives(eval_pos_df, ev_rev_df, ev_man_df)

            train_df = self._add_prompt_columns(train_df)
            eval_df = self._add_prompt_columns(eval_df)

        else:
            if "label" not in pairs_df.columns or pairs_df["label"].nunique() == 1:
                positives_df = pairs_df[pairs_df.get("label", True)][["parent", "child"]].copy()
                pos_rel = self._relations_from_df(positives_df)

                tr_rel, ev_rel = taxonomy_split(
                    pos_rel,
                    train_terms=None,
                    test_size=self._eval_fraction,
                    random_state=self.random_state,
                    verbose=False,
                )
                train_pos_df = self._df_from_relations(tr_rel, label=True)
                eval_pos_df = self._df_from_relations(ev_rel, label=True)

                tr_rev_df, tr_man_df = self._make_negatives(train_pos_df)
                ev_rev_df, ev_man_df = self._make_negatives(eval_pos_df)

                train_df = self._balance_with_negatives(train_pos_df, tr_rev_df, tr_man_df)
                eval_df = self._balance_with_negatives(eval_pos_df, ev_rev_df, ev_man_df)

                train_df = self._add_prompt_columns(train_df)
                eval_df = self._add_prompt_columns(eval_df)

            else:
                positives_df = pairs_df[pairs_df["label"]][["parent", "child"]].copy()
                pos_rel = self._relations_from_df(positives_df)

                tr_rel, ev_rel = taxonomy_split(
                    pos_rel,
                    train_terms=None,
                    test_size=self._eval_fraction,
                    random_state=self.random_state,
                    verbose=False,
                )
                train_pos_df = self._df_from_relations(tr_rel, label=True)
                eval_pos_df = self._df_from_relations(ev_rel, label=True)

                negatives_df = pairs_df[pairs_df["label"]][["parent", "child"]].copy()
                negatives_df = negatives_df.sample(frac=1.0, random_state=self.random_state).reset_index(drop=True)

                n_eval_neg = max(1, int(len(negatives_df) * self._eval_fraction)) if len(negatives_df) > 0 else 0

                eval_neg_df = negatives_df.iloc[:n_eval_neg].copy() if n_eval_neg > 0 else negatives_df.iloc[:0].copy()
                train_neg_df = negatives_df.iloc[n_eval_neg:].copy()

                train_neg_df["label"] = False
                eval_neg_df["label"] = False

                train_df = pd.concat([train_pos_df, train_neg_df], ignore_index=True)
                eval_df = pd.concat([eval_pos_df, eval_neg_df], ignore_index=True)

                train_df = self._add_prompt_columns(train_df)
                eval_df = self._add_prompt_columns(eval_df)

        # Ensure labels are int64
        train_df["label"] = train_df["label"].astype("int64")
        eval_df["label"] = eval_df["label"].astype("int64")

        # Sequential fine-tuning across prompts
        for i in range(self.n_prompts):
            prompt_col = f"prompt_{i + 1}"
            train_ds = Dataset.from_pandas(train_df[[prompt_col, "label"]].reset_index(drop=True))
            eval_ds = Dataset.from_pandas(eval_df[[prompt_col, "label"]].reset_index(drop=True))

            train_ds = train_ds.rename_column("label", "labels")
            eval_ds = eval_ds.rename_column("label", "labels")

            def tokenize_batch(batch):
                """Tokenize a batch for the current prompt column with truncation/padding."""
                return self.tokenizer(batch[prompt_col], padding="max_length", truncation=True)

            train_ds = train_ds.map(tokenize_batch, batched=True, remove_columns=[prompt_col])
            eval_ds = eval_ds.map(tokenize_batch, batched=True, remove_columns=[prompt_col])

            train_ds.set_format(type="torch", columns=["input_ids", "attention_mask", "labels"])
            eval_ds.set_format(type="torch", columns=["input_ids", "attention_mask", "labels"])

            trainer = Trainer(
                model=self.model,
                args=self.training_args,
                train_dataset=train_ds,
                eval_dataset=eval_ds,
            )
            trainer.train()

        self._last_train = train_df
        self._last_eval = eval_df
        return None

class SKHNLPZSLearner(AutoLearner):
    """
    Zero-shot taxonomy learner using an instruction-tuned causal LLM.

    Behavior
    --------
    - Builds a fixed classification prompt listing 9 GeoNames parent classes.
    - For each input row (child term), generates a short completion and parses
      the predicted class from a strict '#[ ... ]#' format.
    - Optionally normalizes the raw prediction to one of the valid 9 labels via:
        * "none"        : keep the parsed text as-is
        * "substring"   : snap to a label if either is a substring of the other
        * "levenshtein" : snap to the closest label by edit distance
        * "auto"        : substring, then Levenshtein if needed
    - Saves raw and normalized predictions to CSV if `save_path` is provided.

    Inputs the learner accepts (via `_to_dataframe`):
    - pandas.DataFrame with columns: ['child', 'parent'] or ['child', 'parent', 'label']
    - list[dict] with keys: 'child', 'parent' (and optionally 'label')
    - list of tuples/lists: (child, parent) or (child, parent, label)
    - OntoLearner-style object exposing .type_taxonomies.taxonomies iterable with (child, parent)
    """

    # Fixed class inventory (GeoNames parents)
    CLASS_LIST = [
        "city, village",
        "country, state, region",
        "forest, heath",
        "mountain, hill, rock",
        "parks, area",
        "road, railroad",
        "spot, building, farm",
        "stream, lake",
        "undersea",
    ]

    # Strict format: #[ ... ]#
    _PREDICTION_PATTERN = re.compile(r"#\[\s*([^\]]+?)\s*\]#")

    def __init__(
        self,
        model_name: str = "Qwen/Qwen2.5-0.5B-Instruct",
        device: Optional[str] = None,  # "cuda" | "cpu" | None (auto)
        max_new_tokens: int = 16,
        save_path: Optional[str] = None,  # directory or full path
        verbose: bool = True,
        normalize_mode: str = "none",  # "none" | "substring" | "levenshtein" | "auto"
        random_state: int = 1403,
    ) -> None:
        """Configure the zero-shot learner.

        Args:
            model_name: HF model id/path for the instruction-tuned causal LLM.
            device: Force device ('cuda' or 'cpu'), else auto-detect.
            max_new_tokens: Generation length budget for each completion.
            save_path: Optional CSV path or directory for saving predictions.
            verbose: If True, print progress messages.
            normalize_mode: Post-processing for class names
                ('none' | 'substring' | 'levenshtein' | 'auto').
            random_state: RNG seed for any sampling steps.
        """
        super().__init__()
        self.model_name = model_name
        self.verbose = verbose
        self.max_new_tokens = max_new_tokens
        self.save_path = save_path
        self.normalize_mode = (normalize_mode or "none").lower().strip()
        self.random_state = random_state

        random.seed(self.random_state)

        # Device: auto-detect CUDA if not specified
        if device is None:
            self._has_cuda = torch.cuda.is_available()
        else:
            self._has_cuda = device == "cuda"
        self._pipe_device = 0 if self._has_cuda else -1
        self._model_device_map = {"": "cuda"} if self._has_cuda else None

        self._tokenizer = None
        self._model = None
        self._pipeline = None

        # Prompt template used for every example
        self._classification_prompt = (
            "My task is classification. My classes are as follows: "
            "(city, village), (country, state, region), (forest, heath), "
            "(mountain, hill, rock), (parks, area), (road, railroad), "
            "(spot, building, farm), (stream, lake), (undersea). "
            'I will provide you with a phrase like "wadi mouth". '
            "The name of each class is placed within a pair of parentheses. "
            "I want you to choose the most appropriate class from those mentioned above "
            "based on the given phrase and present it in a format like #[parks, area]#. "
            "So, the general format for each response will be #[class name]#. "
            "Pay attention to the format of the response. Start with a '#' character, "
            "include the class name inside it, and end with another '#' character. "
            "Additionally, make sure to include a '#' character at the end to indicate "
            "that the answer is complete. I don't need any additional explanations."
        )

    def load(self, model_id: str = "") -> None:
        """
        Load tokenizer, model, and text-generation pipeline.

        Args:
            model_id: Optional HF id/path override; defaults to `self.model_name`.

        Side Effects:
            Initializes the tokenizer and model, configures the generation
            pipeline on CPU/GPU, and sets a pad token if absent.
        """
        model_id = model_id or self.model_name
        if self.verbose:
            print(f"[ZeroShotTaxonomyLearner] Loading {model_id}")

        self._tokenizer = AutoTokenizer.from_pretrained(model_id)

        # Ensure a pad token is set for generation
        if self._tokenizer.pad_token_id is None and self._tokenizer.eos_token_id is not None:
            self._tokenizer.pad_token = self._tokenizer.eos_token

        self._model = AutoModelForCausalLM.from_pretrained(
            model_id,
            device_map=self._model_device_map,
            torch_dtype="auto",
        )

        self._pipeline = pipeline(
            task="text-generation",
            model=self._model,
            tokenizer=self._tokenizer,
            device=self._pipe_device,  # 0 for GPU, -1 for CPU
        )

        if self.verbose:
            print("Device set to use", "cuda" if self._has_cuda else "cpu")
            print("[ZeroShotTaxonomyLearner] Model loaded.")

    def _taxonomy_discovery(self, data: Any, test: bool = False) -> Optional[List[Dict[str, str]]]:
        """
        Zero-shot prediction over all incoming rows (no filtering/augmentation).

        Args:
            data: One of {DataFrame, list[dict], list[tuple], Ontology-like}.
            test: If False, training is skipped (zero-shot learner), and None is returned.

        Returns:
            On `test=True`, a list of dicts [{'parent': predicted_label, 'child': child}, ...].
            On `test=False`, returns None.
        """
        if not test:
            if self.verbose:
                print("[ZeroShot] Training skipped (zero-shot).")
            return None

        df = self._to_dataframe(data)

        if self.verbose:
            print(f"[ZeroShot] Incoming rows: {len(df)}; columns: {list(df.columns)}")

        eval_df = pd.DataFrame(df).reset_index(drop=True)
        if eval_df.empty:
            return []

        # Prepare columns for inspection and saving
        eval_df["prediction_raw"] = ""
        eval_df["prediction_sub"] = ""
        eval_df["prediction_lvn"] = ""
        eval_df["prediction_auto"] = ""
        eval_df["prediction"] = ""  # final (per normalize_mode)

        # Generate predictions row by row
        for idx, row in eval_df.iterrows():
            child_term = str(row["child"])
            raw_text, parsed_raw = self._generate_and_parse(child_term)

            # Choose a string to normalize (parsed token if matched, otherwise whole output)
            basis = parsed_raw if parsed_raw != "unknown" else raw_text

            # Compute all normalization variants
            sub_norm = self._normalize_substring_only(basis)
            lvn_norm = self._normalize_levenshtein_only(basis)
            auto_norm = self._normalize_auto(basis)

            # Final selection by mode
            if self.normalize_mode == "none":
                final_label = parsed_raw
            elif self.normalize_mode == "substring":
                final_label = sub_norm
            elif self.normalize_mode == "levenshtein":
                final_label = lvn_norm
            elif self.normalize_mode == "auto":
                final_label = auto_norm
            else:
                final_label = parsed_raw  # fallback

            # Persist to DataFrame for inspection/export
            eval_df.at[idx, "prediction_raw"] = parsed_raw
            eval_df.at[idx, "prediction_sub"] = sub_norm
            eval_df.at[idx, "prediction_lvn"] = lvn_norm
            eval_df.at[idx, "prediction_auto"] = auto_norm
            eval_df.at[idx, "prediction"] = final_label

        # Return in the format expected by the pipeline
        return [
            {"parent": p, "child": c}
            for p, c in zip(eval_df["prediction"], eval_df["child"])
        ]

    def _generate_and_parse(self, child_term: str) -> (str, str):
        """
        Generate a completion for the given child term and extract the raw predicted class
        using the strict '#[ ... ]#' pattern.

        Args:
            child_term: The child label to classify into one of the fixed classes.

        Returns:
            Tuple `(raw_generation_text, parsed_prediction_or_unknown)`, where the second
            element is either the text inside '#[ ... ]#' or the string 'unknown'.
        """
        messages = [
            {"role": "system", "content": "You are a helpful classifier."},
            {"role": "user", "content": f"{self._classification_prompt} {child_term}"},
        ]

        prompt = self._tokenizer.apply_chat_template(
            messages,
            tokenize=False,
            add_generation_prompt=True,
        )

        generation = self._pipeline(prompt,
                                    max_new_tokens=self.max_new_tokens,
                                    do_sample=False,
                                    temperature=0.0,
                                    top_p=1.0,
                                    eos_token_id=self._tokenizer.eos_token_id,
                                    pad_token_id=self._tokenizer.eos_token_id,
                                    return_full_text=False)[0]["generated_text"]

        match = self._PREDICTION_PATTERN.search(generation)
        parsed = match.group(1).strip() if match else "unknown"
        return generation, parsed

    def _normalize_substring_only(self, text: str) -> str:
        """
        Snap to a label if the string is equal to / contained in / contains a valid label (case-insensitive).

        Args:
            text: Raw class text to normalize.

        Returns:
            One of `CLASS_LIST` on a match; otherwise 'unknown'.
        """
        if not isinstance(text, str):
            return "unknown"
        lowered = text.strip().lower()
        if not lowered:
            return "unknown"

        for label in self.CLASS_LIST:
            label_lower = label.lower()
            if lowered == label_lower or lowered in label_lower or label_lower in lowered:
                return label
        return "unknown"

    def _normalize_levenshtein_only(self, text: str) -> str:
        """
        Snap to the nearest label by Levenshtein (edit) distance.

        Args:
            text: Raw class text to normalize.

        Returns:
            The nearest label in `CLASS_LIST`, or 'unknown' if input is empty/invalid.
        """
        if not isinstance(text, str):
            return "unknown"
        lowered = text.strip().lower()
        if not lowered:
            return "unknown"

        best_label = None
        best_distance = 10**9
        for label in self.CLASS_LIST:
            label_lower = label.lower()
            distance = Levenshtein.distance(lowered, label_lower)
            if distance < best_distance:
                best_distance = distance
                best_label = label
        return best_label or "unknown"

    def _normalize_auto(self, text: str) -> str:
        """
        Cascade: try substring-first; if no match, fall back to Levenshtein snapping.

        Args:
            text: Raw class text to normalize.

        Returns:
            Normalized label string or 'unknown'.
        """
        snapped = self._normalize_substring_only(text)
        return snapped if snapped != "unknown" else self._normalize_levenshtein_only(text)

    def _to_dataframe(self, data: Any) -> pd.DataFrame:
        """
        Normalize various input formats into a DataFrame.

        Supported inputs:
            * pandas.DataFrame with columns ['child','parent',('label')]
            * list[dict] with keys 'child','parent',('label')
            * list of tuples/lists: (child, parent) or (child, parent, label)
            * Ontology-like object with `.type_taxonomies.taxonomies`

        Args:
            data: The source object to normalize.

        Returns:
            A pandas DataFrame with standardized columns.

        Raises:
            ValueError: If the input type/shape is not recognized.
        """
        if isinstance(data, pd.DataFrame):
            df = data.copy()
            df.columns = [str(c).lower() for c in df.columns]
            return df.reset_index(drop=True)

        if isinstance(data, list) and data and isinstance(data[0], dict):
            rows = [{str(k).lower(): v for k, v in d.items()} for d in data]
            return pd.DataFrame(rows).reset_index(drop=True)

        if isinstance(data, (list, tuple)) and data:
            first = data[0]
            if isinstance(first, (list, tuple)) and not isinstance(first, dict):
                n = len(first)
                if n >= 3:
                    return pd.DataFrame(data, columns=["child", "parent", "label"]).reset_index(drop=True)
                if n == 2:
                    return pd.DataFrame(data, columns=["child", "parent"]).reset_index(drop=True)

        try:
            type_taxonomies = getattr(data, "type_taxonomies", None)
            if type_taxonomies is not None:
                taxonomies = getattr(type_taxonomies, "taxonomies", None)
                if taxonomies is not None:
                    rows = []
                    for rel in taxonomies:
                        parent = getattr(rel, "parent", None)
                        child = getattr(rel, "child", None)
                        label = getattr(rel, "label", None) if hasattr(rel, "label") else None
                        if parent is not None and child is not None:
                            rows.append({"child": child, "parent": parent, "label": label})
                    if rows:
                        return pd.DataFrame(rows).reset_index(drop=True)
        except Exception:
            pass

        raise ValueError(
            "Unsupported data format. Provide a DataFrame, a list of dicts, "
            "a list of (child, parent[, label]) tuples/lists, or an object with "
            ".type_taxonomies.taxonomies."
        )

    def _resolve_save_path(self, save_path: str, default_filename: str) -> str:
        """
        Resolve a target file path from a directory or path-like input.

        If `save_path` points to a directory, joins it with `default_filename`.
        If it already looks like a file path (has an extension), returns as-is.

        Args:
            save_path: Directory or file path supplied by the caller.
            default_filename: Basename to use when `save_path` is a directory.

        Returns:
            A concrete file path where outputs can be written.
        """
        base = os.path.basename(save_path)
        has_ext = os.path.splitext(base)[1] != ""
        return save_path if has_ext else os.path.join(save_path, default_filename)
