"""
PyTorch 모델을 TensorFlow Lite 형식으로 변환하는 스크립트
"""

import torch
import torch.nn as nn
import numpy as np
import os

# 선택적 임포트
ONNX_AVAILABLE = False
TF_AVAILABLE = False
ONNX_TF_AVAILABLE = False

try:
    import onnx
    ONNX_AVAILABLE = True
except (ImportError, SyntaxError, Exception) as e:
    ONNX_AVAILABLE = False
    if not isinstance(e, ImportError):
        print(f"⚠️  onnx 패키지 로드 중 오류 발생: {type(e).__name__}")

try:
    import tensorflow as tf
    TF_AVAILABLE = True
except (ImportError, SyntaxError, Exception) as e:
    TF_AVAILABLE = False
    if not isinstance(e, ImportError):
        print(f"⚠️  tensorflow 패키지 로드 중 오류 발생: {type(e).__name__}")

try:
    # onnx-tf는 실제로 사용할 때 임포트하도록 변경
    # from onnx_tf.backend import prepare
    ONNX_TF_AVAILABLE = True
except (ImportError, SyntaxError, Exception) as e:
    ONNX_TF_AVAILABLE = False
    if not isinstance(e, ImportError):
        print(f"⚠️  onnx-tf 패키지 로드 중 오류 발생: {type(e).__name__}")


class FatigueNet(nn.Module):
    """CNN + GRU 기반 피로도 예측 모델 (PyTorch 버전)"""
    
    def __init__(self, input_dim=2, hidden_dim=64, num_layers=2, output_dim=1):
        super(FatigueNet, self).__init__()
        
        # CNN 부분
        self.conv1 = nn.Conv1d(
            in_channels=input_dim,
            out_channels=32,
            kernel_size=1,
            padding=0
        )
        self.conv2 = nn.Conv1d(
            in_channels=32,
            out_channels=64,
            kernel_size=1,
            padding=0
        )
        self.relu = nn.ReLU()
        
        # GRU 부분 (TFLite 호환성을 위해 linear_before_reset=False)
        self.gru = nn.GRU(
            input_size=64,
            hidden_size=hidden_dim,
            num_layers=num_layers,
            batch_first=True,
            dropout=0.2 if num_layers > 1 else 0
        )
        
        # Fully Connected 레이어
        self.fc = nn.Linear(hidden_dim, output_dim)
        self.dropout = nn.Dropout(0.3)
    
    def forward(self, x):
        if x.dim() == 2:
            x = x.unsqueeze(1)
        
        x = x.permute(0, 2, 1)
        
        x = self.conv1(x)
        x = self.relu(x)
        
        x = self.conv2(x)
        x = self.relu(x)
        
        x = x.permute(0, 2, 1)
        
        gru_out, _ = self.gru(x)
        
        last_output = gru_out[:, -1, :]
        
        last_output = self.dropout(last_output)
        output = self.fc(last_output)
        
        return output


def convert_pytorch_to_tflite(
    pytorch_model_path='./model/fatigue_net_v2.pt',
    tflite_model_path='./model/fatigue_net_v2.tflite',
    input_shape=(1, 1, 2)  # (batch, seq_len, features)
):
    """
    PyTorch 모델을 TensorFlow Lite로 변환
    
    Args:
        pytorch_model_path: PyTorch 모델 파일 경로
        tflite_model_path: 저장할 TFLite 모델 파일 경로
        input_shape: 입력 텐서 형태 (batch, seq_len, features)
    """
    print("=" * 80)
    print("PyTorch 모델을 TensorFlow Lite로 변환")
    print("=" * 80)
    
    # 필수 패키지 확인
    if not ONNX_AVAILABLE or not TF_AVAILABLE or not ONNX_TF_AVAILABLE:
        print("\n❌ 필수 패키지가 설치되지 않았거나 호환성 문제가 있습니다.")
        print("\n📋 Python 버전 확인:")
        import sys
        print(f"   현재 Python 버전: {sys.version}")
        print(f"   권장 Python 버전: 3.10 이상")
        
        if sys.version_info < (3, 10):
            print("\n⚠️  Python 3.9에서는 일부 패키지 호환성 문제가 있을 수 있습니다.")
            print("   Python 3.10 이상으로 업그레이드하거나, 다음을 시도하세요:")
            print("   - 가상환경에서 Python 3.10+ 사용")
            print("   - 또는 호환되는 패키지 버전 설치")
        
        print("\n📦 설치 명령어:")
        print("  권장 버전 (Python 3.10 이상):")
        print("    pip install onnx==1.15.0 onnx-tf==1.10.0 tensorflow==2.15.0")
        print("\n⚠️  참고: Python 3.9에서는 일부 패키지 설치 중 에러가 발생할 수 있습니다.")
        print("   Python 3.10 이상 사용을 강력히 권장합니다.")
        
        print("\n❌ TFLite 변환은 필수입니다. 모바일 디바이스에서 실행하기 위해 필요합니다.")
        print("   필수 패키지를 설치하고 다시 시도하세요.")
        return False
    
    # 1️⃣ PyTorch 모델 로드
    print("\n1️⃣ PyTorch 모델 로드 중...")
    if not os.path.exists(pytorch_model_path):
        raise FileNotFoundError(f"모델 파일을 찾을 수 없습니다: {pytorch_model_path}")
    
    # TorchScript 파일인지 먼저 확인
    try:
        checkpoint = torch.jit.load(pytorch_model_path, map_location='cpu')
        if isinstance(checkpoint, torch.jit.ScriptModule):
            raise ValueError(
                f"❌ {pytorch_model_path}는 TorchScript 형식입니다.\n"
                "TFLite 변환을 위해서는 PyTorch state_dict 형식 모델이 필요합니다.\n"
                "모델을 다시 학습하거나 올바른 형식의 모델 파일을 사용하세요."
            )
    except:
        pass
    
    # 일반 PyTorch 모델 로드
    checkpoint = torch.load(pytorch_model_path, map_location='cpu')
    
    # 일반 PyTorch 모델인지 확인
    if not isinstance(checkpoint, dict) or 'model_state_dict' not in checkpoint:
        raise ValueError(
            f"❌ 올바른 PyTorch 모델 형식이 아닙니다.\n"
            f"'{pytorch_model_path}' 파일에 'model_state_dict' 키가 필요합니다.\n"
            "모델을 다시 학습하거나 올바른 형식의 모델 파일을 사용하세요."
        )
    
    model_config = checkpoint.get('model_config', {
        'input_dim': 2,
        'hidden_dim': 64,
        'num_layers': 2,
        'output_dim': 1
    })
    
    model = FatigueNet(**model_config)
    model.load_state_dict(checkpoint['model_state_dict'])
    model.eval()
    print(f"✅ 모델 로드 완료: {pytorch_model_path}")
    print(f"   모델 설정: {model_config}\n")
    
    # 2️⃣ ONNX로 변환
    print("2️⃣ ONNX 형식으로 변환 중...")
    onnx_model_path = './model/fatigue_net_v2.onnx'
    os.makedirs('./model', exist_ok=True)
    
    # 더미 입력 생성 (고정 batch_size=1로 TFLite 호환성 향상)
    dummy_input = torch.randn(1, 1, 2)  # (batch=1, seq_len=1, features=2)
    
    try:
        # GRU를 RNN으로 변환하거나 TFLite 호환 옵션 사용
        torch.onnx.export(
            model,
            dummy_input,
            onnx_model_path,
            export_params=True,
            opset_version=11,  # onnx-tf 호환성을 위해 11로 낮춤
            do_constant_folding=True,
            input_names=['input'],
            output_names=['output'],
            dynamic_axes={
                'input': {0: 'batch_size', 1: 'sequence_length'},
                'output': {0: 'batch_size'}
            },
            # GRU 관련 호환성 옵션
            custom_opsets=None,
            verbose=False
        )
        print(f"✅ ONNX 변환 완료: {onnx_model_path}\n")
    except Exception as e:
        print(f"⚠️  ONNX 변환 중 경고 (계속 진행): {e}\n")
    
    # 3️⃣ ONNX를 TensorFlow로 변환
    print("3️⃣ TensorFlow 형식으로 변환 중...")
    try:
        from onnx_tf.backend import prepare
        
        # ONNX 모델 로드 및 GRU 속성 수정
        onnx_model = onnx.load(onnx_model_path)
        
        # GRU 노드의 linear_before_reset 속성을 0으로 설정 (TensorFlow 호환)
        for node in onnx_model.graph.node:
            if node.op_type == 'GRU':
                # linear_before_reset 속성을 찾아서 0으로 설정
                for attr in node.attribute:
                    if attr.name == 'linear_before_reset':
                        attr.i = 0
                        break
                else:
                    # linear_before_reset 속성이 없으면 추가
                    attr = onnx.helper.make_attribute('linear_before_reset', 0)
                    node.attribute.append(attr)
        
        tf_rep = prepare(onnx_model)
        
        # TensorFlow SavedModel로 저장
        tf_model_path = './model/tf_model'
        tf_rep.export_graph(tf_model_path)
        print(f"✅ TensorFlow 변환 완료: {tf_model_path}\n")
    except Exception as e:
        print(f"❌ TensorFlow 변환 실패: {e}")
        print("⚠️  ONNX-TF 변환이 실패했습니다.\n")
        print("❌ TFLite 변환은 필수입니다. 모바일 디바이스에서 실행하기 위해 필요합니다.")
        print("   에러를 해결하고 다시 시도하세요.")
        return False
    
    # 4️⃣ TensorFlow Lite로 변환
    print("4️⃣ TensorFlow Lite 형식으로 변환 중...")
    
    # TensorFlow Lite 변환기 생성
    converter = tf.lite.TFLiteConverter.from_saved_model(tf_model_path)
    
    # GRU 등 복잡한 연산을 위한 설정
    converter.target_spec.supported_ops = [
        tf.lite.OpsSet.TFLITE_BUILTINS,
        tf.lite.OpsSet.SELECT_TF_OPS
    ]
    converter._experimental_lower_tensor_list_ops = False
    
    # 최적화 옵션 설정 (선택사항)
    converter.optimizations = [tf.lite.Optimize.DEFAULT]
    
    # 변환 실행
    tflite_model = converter.convert()
    
    # TFLite 모델 저장
    with open(tflite_model_path, 'wb') as f:
        f.write(tflite_model)
    
    print(f"✅ TensorFlow Lite 변환 완료: {tflite_model_path}")
    
    # 모델 크기 확인
    model_size = os.path.getsize(tflite_model_path) / (1024 * 1024)  # MB
    print(f"   모델 크기: {model_size:.2f} MB\n")
    
    # 5️⃣ 변환된 모델 테스트
    print("5️⃣ 변환된 모델 테스트 중...")
    try:
        interpreter = tf.lite.Interpreter(model_path=tflite_model_path)
        interpreter.allocate_tensors()
        
        input_details = interpreter.get_input_details()
        output_details = interpreter.get_output_details()
        
        print(f"   입력 형태: {input_details[0]['shape']}")
        print(f"   출력 형태: {output_details[0]['shape']}")
        
        # 테스트 입력 (고정 크기)
        test_input = np.random.randn(1, 1, 2).astype(np.float32)
        interpreter.set_tensor(input_details[0]['index'], test_input)
        interpreter.invoke()
        test_output = interpreter.get_tensor(output_details[0]['index'])
        
        print(f"   테스트 출력: {test_output[0][0]:.4f}")
        print("   ✅ 모델 테스트 성공\n")
    except Exception as e:
        print(f"   ⚠️  모델 테스트 중 경고: {e}")
        print("   (모델은 생성되었지만 테스트는 실패했습니다. 모바일 디바이스에서 Flex ops가 필요할 수 있습니다.)\n")
    
    # 중간 파일 정리 (선택사항)
    print("6️⃣ 중간 파일 정리 중...")
    try:
        os.remove(onnx_model_path)
        import shutil
        shutil.rmtree(tf_model_path)
        print("✅ 중간 파일 정리 완료\n")
    except Exception as e:
        print(f"⚠️  중간 파일 정리 실패 (무시 가능): {e}\n")
    
    print("=" * 80)
    print(f"✅ 변환 완료!")
    print(f"   TFLite 모델: {tflite_model_path}")
    print("=" * 80)
    return True


def main():
    """메인 함수"""
    try:
        success = convert_pytorch_to_tflite(
            pytorch_model_path='./model/fatigue_net_v2.pt',
            tflite_model_path='./model/fatigue_net_v2.tflite'
        )
        if not success:
            return 1
    except Exception as e:
        print(f"\n❌ 변환 실패: {e}")
        import traceback
        traceback.print_exc()
        return 1
    
    return 0


if __name__ == "__main__":
    exit(main())