"""
Part 7 딥러닝 기초 테스트 코드
"""

import os
import sys
import unittest

import numpy as np

# PyTorch 관련 import
try:
    import torch
    import torch.nn as nn
    import torch.optim as optim
    from torch.utils.data import DataLoader, TensorDataset

    TORCH_AVAILABLE = True
except ImportError:
    TORCH_AVAILABLE = False

# Scikit-learn import
from sklearn.datasets import load_iris
from sklearn.metrics import accuracy_score
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler

# 상위 디렉토리의 모듈을 import하기 위한 경로 추가
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))


class TestDeepLearningBasics(unittest.TestCase):
    """딥러닝 기초 테스트 클래스"""

    def setUp(self):
        """테스트 설정"""
        if not TORCH_AVAILABLE:
            self.skipTest("PyTorch가 설치되지 않았습니다.")

        # 붓꽃 데이터셋 로드
        iris = load_iris()
        X = iris.data
        y = iris.target

        # 훈련/테스트 분할
        X_train, X_test, y_train, y_test = train_test_split(
            X, y, test_size=0.2, random_state=42, stratify=y
        )

        # 스케일링
        scaler = StandardScaler()
        X_train_scaled = scaler.fit_transform(X_train)
        X_test_scaled = scaler.transform(X_test)

        # PyTorch 텐서로 변환
        self.X_train_tensor = torch.FloatTensor(X_train_scaled)
        self.y_train_tensor = torch.LongTensor(y_train)
        self.X_test_tensor = torch.FloatTensor(X_test_scaled)
        self.y_test_tensor = torch.LongTensor(y_test)

        # 모델 파라미터
        self.input_size = X_train.shape[1]
        self.hidden_size = 32
        self.num_classes = len(np.unique(y_train))

    def test_pytorch_installation(self):
        """PyTorch 설치 확인 테스트"""
        self.assertTrue(TORCH_AVAILABLE, "PyTorch가 설치되지 않았습니다.")
        self.assertTrue(hasattr(torch, "__version__"))

    def test_tensor_creation(self):
        """텐서 생성 테스트"""
        if not TORCH_AVAILABLE:
            self.skipTest("PyTorch가 설치되지 않았습니다.")

        # 텐서 크기 확인
        self.assertEqual(self.X_train_tensor.shape, (120, 4))
        self.assertEqual(self.y_train_tensor.shape, (120,))

        # 텐서 타입 확인
        self.assertEqual(self.X_train_tensor.dtype, torch.float32)
        self.assertEqual(self.y_train_tensor.dtype, torch.int64)

    def test_simple_neural_network(self):
        """간단한 신경망 테스트"""
        if not TORCH_AVAILABLE:
            self.skipTest("PyTorch가 설치되지 않았습니다.")

        # 간단한 신경망 클래스 정의
        class SimpleNN(nn.Module):
            def __init__(self, input_size, hidden_size, num_classes):
                super(SimpleNN, self).__init__()
                self.fc1 = nn.Linear(input_size, hidden_size)
                self.relu = nn.ReLU()
                self.fc2 = nn.Linear(hidden_size, num_classes)

            def forward(self, x):
                out = self.fc1(x)
                out = self.relu(out)
                out = self.fc2(out)
                return out

        # 모델 생성
        model = SimpleNN(self.input_size, self.hidden_size, self.num_classes)

        # 모델 구조 확인
        self.assertEqual(
            len(list(model.parameters())), 4
        )  # fc1.weight, fc1.bias, fc2.weight, fc2.bias

        # 순전파 테스트
        with torch.no_grad():
            output = model(self.X_test_tensor)
            self.assertEqual(output.shape, (30, 3))  # 30개 샘플, 3개 클래스

    def test_loss_function(self):
        """손실 함수 테스트"""
        if not TORCH_AVAILABLE:
            self.skipTest("PyTorch가 설치되지 않았습니다.")

        # 크로스 엔트로피 손실 함수
        criterion = nn.CrossEntropyLoss()

        # 더미 출력과 타겟 생성
        dummy_output = torch.randn(5, 3)  # 5개 샘플, 3개 클래스
        dummy_target = torch.LongTensor([0, 1, 2, 1, 0])

        # 손실 계산
        loss = criterion(dummy_output, dummy_target)

        # 손실이 스칼라인지 확인
        self.assertEqual(loss.dim(), 0)
        self.assertGreater(loss.item(), 0)

    def test_optimizer(self):
        """옵티마이저 테스트"""
        if not TORCH_AVAILABLE:
            self.skipTest("PyTorch가 설치되지 않았습니다.")

        # 간단한 모델 생성
        model = nn.Linear(4, 3)

        # Adam 옵티마이저
        optimizer = optim.Adam(model.parameters(), lr=0.01)

        # 옵티마이저 속성 확인
        self.assertEqual(len(optimizer.param_groups), 1)
        self.assertEqual(optimizer.param_groups[0]["lr"], 0.01)

    def test_dataloader(self):
        """DataLoader 테스트"""
        if not TORCH_AVAILABLE:
            self.skipTest("PyTorch가 설치되지 않았습니다.")

        # 데이터셋 생성
        dataset = TensorDataset(self.X_train_tensor, self.y_train_tensor)
        dataloader = DataLoader(dataset, batch_size=16, shuffle=True)

        # DataLoader 속성 확인
        self.assertEqual(len(dataset), 120)
        self.assertEqual(len(dataloader), 8)  # 120 / 16 = 7.5 -> 8 배치

        # 첫 번째 배치 확인
        for features, labels in dataloader:
            self.assertEqual(features.shape[1], 4)  # 특성 수
            self.assertLessEqual(features.shape[0], 16)  # 배치 크기
            break

    def test_model_training_step(self):
        """모델 훈련 단계 테스트"""
        if not TORCH_AVAILABLE:
            self.skipTest("PyTorch가 설치되지 않았습니다.")

        # 간단한 신경망 클래스 정의
        class SimpleNN(nn.Module):
            def __init__(self, input_size, hidden_size, num_classes):
                super(SimpleNN, self).__init__()
                self.fc1 = nn.Linear(input_size, hidden_size)
                self.relu = nn.ReLU()
                self.fc2 = nn.Linear(hidden_size, num_classes)

            def forward(self, x):
                out = self.fc1(x)
                out = self.relu(out)
                out = self.fc2(out)
                return out

        # 모델, 손실함수, 옵티마이저 생성
        model = SimpleNN(self.input_size, self.hidden_size, self.num_classes)
        criterion = nn.CrossEntropyLoss()
        optimizer = optim.Adam(model.parameters(), lr=0.01)

        # 초기 가중치 저장
        initial_weights = list(model.parameters())[0].clone()

        # 한 번의 훈련 스텝
        model.train()
        outputs = model(self.X_train_tensor)
        loss = criterion(outputs, self.y_train_tensor)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        # 가중치가 업데이트되었는지 확인
        updated_weights = list(model.parameters())[0]
        self.assertFalse(torch.equal(initial_weights, updated_weights))

    def test_model_evaluation(self):
        """모델 평가 테스트"""
        if not TORCH_AVAILABLE:
            self.skipTest("PyTorch가 설치되지 않았습니다.")

        # 간단한 모델 생성 및 훈련
        model = nn.Linear(self.input_size, self.num_classes)
        criterion = nn.CrossEntropyLoss()
        optimizer = optim.Adam(model.parameters(), lr=0.1)

        # 간단한 훈련 (몇 번의 에포크)
        model.train()
        for epoch in range(10):
            outputs = model(self.X_train_tensor)
            loss = criterion(outputs, self.y_train_tensor)
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

        # 평가 모드로 전환
        model.eval()
        with torch.no_grad():
            test_outputs = model(self.X_test_tensor)
            _, predicted = torch.max(test_outputs.data, 1)

        # 예측 결과 확인
        self.assertEqual(len(predicted), len(self.y_test_tensor))

        # 정확도 계산 (너무 낮지 않은지 확인)
        accuracy = accuracy_score(self.y_test_tensor.numpy(), predicted.numpy())
        self.assertGreaterEqual(accuracy, 0.2)  # 랜덤보다는 높은 성능


class TestTensorOperations(unittest.TestCase):
    """텐서 연산 테스트 클래스"""

    def test_tensor_basic_operations(self):
        """기본 텐서 연산 테스트"""
        if not TORCH_AVAILABLE:
            self.skipTest("PyTorch가 설치되지 않았습니다.")

        # 텐서 생성
        a = torch.tensor([[1, 2], [3, 4]], dtype=torch.float32)
        b = torch.tensor([[5, 6], [7, 8]], dtype=torch.float32)

        # 덧셈
        c = a + b
        expected = torch.tensor([[6, 8], [10, 12]], dtype=torch.float32)
        self.assertTrue(torch.equal(c, expected))

        # 행렬 곱셈
        d = torch.mm(a, b)
        expected_mm = torch.tensor([[19, 22], [43, 50]], dtype=torch.float32)
        self.assertTrue(torch.equal(d, expected_mm))

    def test_tensor_reshaping(self):
        """텐서 reshape 테스트"""
        if not TORCH_AVAILABLE:
            self.skipTest("PyTorch가 설치되지 않았습니다.")

        # 텐서 생성
        x = torch.arange(12)

        # Reshape
        y = x.view(3, 4)
        self.assertEqual(y.shape, (3, 4))

        # Flatten
        z = y.view(-1)
        self.assertEqual(z.shape, (12,))


def run_tests():
    """테스트 실행 함수"""
    # 테스트 스위트 생성
    test_suite = unittest.TestSuite()

    # 테스트 클래스 추가
    test_suite.addTest(unittest.makeSuite(TestDeepLearningBasics))
    test_suite.addTest(unittest.makeSuite(TestTensorOperations))

    # 테스트 실행
    runner = unittest.TextTestRunner(verbosity=2)
    result = runner.run(test_suite)

    # 결과 출력
    print(f"\n=== 테스트 결과 ===")
    print(f"실행된 테스트: {result.testsRun}")
    print(f"실패한 테스트: {len(result.failures)}")
    print(f"오류가 발생한 테스트: {len(result.errors)}")

    if result.failures:
        print("\n실패한 테스트:")
        for test, error in result.failures:
            print(f"- {test}: {error}")

    if result.errors:
        print("\n오류가 발생한 테스트:")
        for test, error in result.errors:
            print(f"- {test}: {error}")

    return result.wasSuccessful()


if __name__ == "__main__":
    # 직접 실행 시 테스트 수행
    success = run_tests()

    if success:
        print("\n모든 테스트가 성공적으로 완료되었습니다! ✅")
    else:
        print("\n일부 테스트가 실패했습니다. ❌")
        sys.exit(1)
