from abc import abstractmethod
from typing import List

import torch

from ..base.learner import BaseLearner 
from ...mid.constructors.concept_graph import ConceptGraph


class IndependentLearner(BaseLearner):
    def __init__(
        self, 
        graph_levels: List[str]=None,
        **kwargs
    ):
        super(IndependentLearner, self).__init__(**kwargs)
        if graph_levels is not None:
            graph_levels = [[var.concepts[0] for var in level] for level in graph_levels]
            self.graph_levels = graph_levels[1:]
            self.roots = self.graph_levels[0]

    @abstractmethod
    def concept_encoder(self):
        """The concept encoder (input -> concepts) 
        that must be defined by subclasses."""
        pass

    @abstractmethod
    def concept_predictor(self):
        """The concept predictor (concepts -> concepts)
        that must be defined by subclasses."""
        pass

    def predict_all_levels(self, inputs, concepts=None, use_ground_truth=True):
        """Predict concepts for all levels.

        Args:
            inputs (dict): Input dictionary containing 'x'.
            concepts (dict, optional): Ground truth concepts. Required if use_ground_truth=True.
            use_ground_truth (bool): If True, use ground truth from previous levels.
                                     If False, use predicted concepts (cascade).
        
        Returns:
            torch.Tensor: Concatenated predictions for all levels.
        """ 
        all_predictions = []
        
        for level in self.graph_levels:
            # Root level: always use input encoder
            if level == self.roots:
                level_out = self.concept_encoder(x=inputs['x'], query=level)
            # Non-root levels: use ground truth or predictions
            else:
                if use_ground_truth:
                    level_evidence = {}
                    for c_name in level:
                        # extract ground truth endogenous (e.g., logits) for concept
                        idx = self.concept_names.index(c_name)
                        c_true = concepts['c'][:, idx]
                        if self.concept_annotations.types[idx] == 'continuous':
                            raise NotImplementedError("Continuous concepts not yet implemented.")
                        else:
                            # convert to one-hot and then to endogenous
                            card = self.concept_annotations.cardinalities[idx]
                            c_one_hot = torch.nn.functional.one_hot(
                                c_true.long(), 
                                num_classes=card
                            ).float()
                            # Apply logit to one-hot (values are 0 or 1)
                            # FIXME: should we done something different for 
                            # other distributions?
                            c_endogenous = torch.logit(c_one_hot, eps=1e-7)
                            level_evidence[c_name] = c_endogenous
                    
                    # Use ground truth concepts (as logits) from previous levels
                    level_out = self.concept_predictor(evidence=level_evidence, query=level)
                else:
                    # Use predicted concepts from previous levels (cascade)
                    # FIXME: evidence should be a dict
                    level_out = self.concept_predictor(evidence=level_out, query=level)
            
            all_predictions.append(level_out)
        
        # Concatenate predictions from all levels
        # FIXME: concatenation need to match the expected order
        torch.cat(all_predictions, dim=1)
        return 

    def shared_step(self, batch, step='train'):
        """Shared logic for train/val/test steps.
        
        Args:
            batch (dict): Batch dictionary from dataloader.
            step (str): One of 'train', 'val', or 'test'.
            
        Returns:
            torch.Tensor: Scalar loss value.
        """
        inputs, concepts, transforms = self.unpack_batch(batch)
        batch_size = batch['inputs']['x'].size(0)
        c = c_loss = concepts['c']

        # Training uses ground truth, validation/test use cascade
        use_ground_truth = (step == 'train')
        out = self.predict_all_levels(inputs, concepts, use_ground_truth)

        # --- Compute loss ---
        if self.loss is not None:
            in_loss_dict = self.filter_output_for_loss(out, c_loss)
            loss = self.loss(**in_loss_dict)
            self.log_loss(step, loss, batch_size=batch_size)

        # --- Update and log metrics ---
        metrics_args = self.filter_output_for_metrics(out, c)
        self.update_and_log_metrics(metrics_args, step, batch_size)

        return loss

    def training_step(self, batch):
        """Training step for independent learning.
        
        Uses ground truth concepts from previous levels.
        
        Args:
            batch (dict): Training batch.
            
        Returns:
            torch.Tensor: Training loss.
        """
        return self.shared_step(batch, step='train')

    def validation_step(self, batch):
        """Validation step with cascading predictions.
        
        Uses predicted concepts from previous levels.
        
        Args:
            batch (dict): Validation batch.
            
        Returns:
            torch.Tensor: Validation loss.
        """
        return self.shared_step(batch, step='val')
    
    def test_step(self, batch):
        """Test step with cascading predictions.
        
        Uses predicted concepts from previous levels.
        
        Args:
            batch (dict): Test batch.
            
        Returns:
            torch.Tensor: Test loss.
        """
        return self.shared_step(batch, step='test')
