# soccernetpro/apis/classification.py

from ..core.utils.config import load_config
from ..core.trainer import Trainer
from ..datasets.builder import build_dataset
from ..models.builder import build_model
import os
import torch

HOME_DIR = os.path.expanduser("~")
BASE_DIR = os.path.dirname(os.path.abspath(__name__))
print(BASE_DIR)
class ClassificationAPI:
    def __init__(self, config=None):
        config_path = os.path.join(HOME_DIR, config)
        # Store config (can be path or dict)
        self.config = load_config(config_path)
        self.trainer = Trainer(self.config)

    def train(self, train_set=None, valid_set=None, pretrained=None):
        if pretrained:
            pretrained = os.path.join(HOME_DIR, pretrained)
            self.model, self.processor, _ = self.trainer.load(pretrained)
        else:
            self.model, self.processor = build_model(self.config, self.trainer.device)
        # Build dataset
        train_set = os.path.join(HOME_DIR, train_set or self.config.DATA.annotations.train)
        
        valid_set = os.path.join(HOME_DIR, valid_set or self.config.DATA.annotations.valid)
            
        train_data = build_dataset(self.config, train_set, self.processor, split="train")
        print(f"Train Dataset length: {len(train_data)}")
        frames= train_data[0]
        print(f"Frames shape: {frames['pixel_values'].shape}")  # 
        print(f"Label: {frames['labels']}")
        
        valid_data = build_dataset(self.config, valid_set, self.processor, split="valid") if valid_set else None
        print(f"Valid Dataset length: {len(valid_data)}")
        frames = valid_data[0]
        print(f"Frames shape: {frames['pixel_values'].shape}")  # 
        print(f"Label: {frames['labels']}")
        self.trainer.train_model(self.model, train_data, valid_data)

        save_path = os.path.join(BASE_DIR, self.config.TRAIN.save_dir, "final_model")
        self.trainer.save(model=self.model, path=save_path, processor=self.processor)


    def infer(self, test_set=None, pretrained=None, output_inference=None):
        if pretrained:
            pretrained = os.path.join(HOME_DIR, pretrained)
            self.model, self.processor, _ = self.trainer.load(pretrained)

        # Build dataset
        test_set = os.path.join(HOME_DIR, test_set or self.config.DATA.annotations.valid)
        test_data = build_dataset(self.config, test_set, self.processor, split="test")
        print(f"Test Dataset length: {len(test_data)}")
        frames = test_data[0]
        print(f"Frames shape: {frames['pixel_values'].shape}")  # 
        print(f"Label: {frames['labels']}")

        # # Run inference
        preds, metrics = self.trainer.infer(test_data)
        print(metrics)
        # # Save inference if requested
        # if output_inference:
        #     self.trainer.save_inference(outputs, output_inference)

        # return outputs
        pass

    def demo(self, input_data):
        """Single video demo"""
        #return self.trainer.demo(self.model, input_data)
        pass
