import sys
import torch
import pandas as pd
import numpy as np
import pickle
from torch.utils.data import DataLoader, TensorDataset
from .dataloader import DataProcessor  # DataProcessor는 그대로 사용
from .model import CNNGRUClassificationModel
import pkg_resources


class Predictor:
    def __init__(self, device, model_path=None):
        """
        모델과 장치를 초기화합니다.
        """
        if model_path is None:
            model_path = pkg_resources.resource_filename('saferx', 'model2/model/final_model_revised.pkl')
        
        self.device = device

        # sys.modules 설정을 load_model 안에서 처리
        self.model = self.load_model(model_path)
        self.model.to(self.device)
        self.model.eval()

    def load_model(self, model_path):
        """
        저장된 pickle 모델을 로드합니다.
        """
        # model2를 saferx.model2로 매핑하여 경로 문제 해결
        import saferx.model2 as model2_module
        sys.modules['model2'] = model2_module
        sys.modules['saferx.model2'] = model2_module  # 추가 매핑

        with open(model_path, 'rb') as f:
            model = pickle.load(f)
        return model

    def preprocess_data(self, data_path):
        """
        데이터를 로드하고 전처리한 후 DataLoader로 반환합니다.
        """
        data = pd.read_csv(data_path)
        data = DataProcessor.preprocess_data(data)
        data = DataProcessor.reset_week_numbers(data)
        data = DataProcessor.transform_target(data)

        max_length = DataProcessor.find_max_sequence_length_by_week(data, self.seq_cols)

        # 데이터를 모델에 맞게 준비하고 텐서로 변환
        results = DataProcessor.prepare_data_for_model_by_week(data, max_length, self.seq_cols, self.target_cols)
        X_tensor, _ = DataProcessor.convert_results_to_tensors(results)
        return DataLoader(TensorDataset(X_tensor), batch_size=16, shuffle=False)

    def predict(self, data_loader):
        """
        데이터를 예측하고 결과를 반환합니다.
        """
        predictions = []
        with torch.no_grad():
            for inputs in data_loader:
                inputs = inputs[0].to(self.device)  # 입력 텐서로 변환
                outputs = self.model(inputs)
                # 확률값을 이진값(0 또는 1)으로 변환
                binary_predictions = (outputs >= 0.5).cpu().numpy()
                predictions.extend(binary_predictions)

        return np.array(predictions)

    def save_predictions(self, predictions, output_path):
        """
        예측 결과를 파일로 저장합니다.
        """
        df = pd.DataFrame(predictions, columns=self.target_cols)
        df.to_csv(output_path, index=False)
        print(f'Predictions saved at {output_path}')
