from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split
from torch.utils.data import DataLoader, Dataset
import torch

class IrisDataset(Dataset):
    def __init__(self, mode="train", random_state=0):
        # mode인자를 통해 train인지 test인지 구분하기 -> 사용하는 데이터 달라짐
        iris = load_iris()
        train_X, test_X, train_y, test_y = train_test_split(
            iris.data,
            iris.target,
            stratify=iris.target,
            test_size=0.2,
            random_state=random_state,
        )

        
        if mode == "train": # mode가 train일때는 train_X, train_y를 올려주기
            self.X = torch.tensor(train_X, dtype = torch.float)
             # feature가 되는 X 값을 PyTorch float 형태로 변환
            self.y = torch.tensor(train_y, dtype = torch.long) 
            # Target이 되는 y 값을 PyTorch long 형태로 변환

        else: # mode가 test일떄는 test_X, test_y를 올려주기
            self.X = torch.tensor(test_X, dtype = torch.float)
            self.y = torch.tensor(test_y, dtype = torch.long)
        

    def __len__(self):
        # len()함수로 입력X길이 반환
        return len(self.X)

    def __getitem__(self, idx):
        # getitem함수에서는 idx에 해당하는 x와 y를 반환
        return self.X[idx], self.y[idx]
       

#배치사이즈는 16으로 설정
batch_size = 16

# train_dataset은 인스턴스로 생성
train_dataset = IrisDataset()
#PyTorch의 DataLoader를 이용해서 train_loader인스턴스 생성
train_loader = DataLoader(train_dataset, batch_size = batch_size, shuffle = True)

#test_dataset도 마찬가지로 인스턴스로 생성
test_dataset = IrisDataset()
test_loader = DataLoader(test_dataset, batch_size = batch_size, shuffle = False)

# 선형 레이어를 정의합니다. 이 레이어는 입력 특성과 출력 특성의 수를 받아서 가중치와 편향을 초기화합니다.
class WithoutNNLinear:
    def __init__(self, in_features, out_features):
        # weight는 randn()으로 가우시안 랜덤으로 초기화하기
        self.weight =  torch.randn(out_features, in_features, requires_grad = True)
        #bias는 zeros()로 0으로 초기화하기
        self.bias = torch.zeros(out_features, requires_grad = True) 
        

    # 입력값 x를 받아 선형변환 수행
    def __call__(self, x):
        # weight과 x를 내적하고, bias를 더해서 Linear 연산을 수행
        return torch.matmul(self.weight,x.T).T + self.bias 

    # weight과 bias를 device에 올리기
    def set_device(self, device):
        self.weight = self.weight.to(device)
        self.bias = self.bias.to(device)

    # 모델의 파라미터[Wiehgt, bias]를 반환하기.
    def parameters(self) -> list:
        return [self.weight, self.bias]

class WithoutNNReLU:
    def __init__(self):
        pass

    def __call__(self, x):
        return torch.relu(x)

class WithoutNNMLP:
    def __init__(self, in_features, hidden_features, out_features):

        # 두개의 선형 layer와 relu를 정의해준다.
        self.linear1 = WithoutNNLinear(in_features,hidden_features) 
        self.relu =   WithoutNNReLU()# ReLU 활성화 함수
        self.linear2 = WithoutNNLinear(hidden_features, out_features)

    # 이 메서드를 통해 forward연산을 한다(linear1통과, relu, linear2 통과)
    def __call__(self, x):
        y = self.linear1(x)
        z = self.relu(y)
        return self.linear2(z)

    # Linear와 relu를 device에 올려준다.
    def set_device(self, device):
        self.linear1.set_device(device)
        self.relu.set_device(device)
        self.linear2.set_device(device)

    # 각 Linear의 weight와 bias를 리스트형태로 반환한다.
    def parameters(self) -> list:
        return [self.linear1.weight, self.linear1.bias, self.linear2.weight, self.linear2.bias]

class WithoutNNCrossEntropyLoss:
    def __init__(self, reduce="mean"):
        self.reduce = reduce

    # output인 예측값과 target인 label사이의 CrossEntropyLoss를 계산함.
    def __call__(self, output, target):
	    #output에 softmax를적용한 후 log를 취함
        y_hat = torch.log_softmax(output, dim = 1)
        #target을 원핫벡터로 만들어주기
        target = torch.nn.functional.one_hot(target, num_classes = output.size(1))
        #CrossEntropy공식 적용
        loss = -(torch.sum(target * y_hat)) / output.size(0)
        return loss

# 학습 데이터셋과 데이터 로더를 정의
batch_size = 16
train_dataset = IrisDataset()
train_loader = DataLoader(train_dataset, batch_size = batch_size, shuffle = True)

# 테스트 데이터셋과 데이터 로더를 정의
test_dataset = IrisDataset()
test_loader = DataLoader(test_dataset, batch_size = batch_size, shuffle = False)

# model, Lossfunction, optimizer, trainloader를 받아서 train하는 함수 정의
def train(model, criterion, optimizer, train_loader) -> float:

    # train단계에서는 미분할 수 있도록 설정
    torch.set_grad_enabled(True)

    running_loss = 0
    for X, y in train_loader:
        optimizer.zero_grad() # 그래디언트를 초기화.
        outputs = model(X) # model돌려서 output계산.
        loss = criterion(outputs, y)  # output이랑 label을 lossfunc에다 넣기.
        loss.backward()  # backprop연산.
        optimizer.step()  # 파라미터를 업데이트.
        running_loss += loss.item()  # 전체loss누적.

    return running_loss / len(train_loader)  # 평균 loss을 반환

def test(model, test_loader, criterion) -> tuple[float, float]:

    running_loss = 0
    correct = 0
    total = 0

    with torch.no_grad(): #Test에서는 미분하지 않도록 설정
        for X, y in test_loader:
            predictions = model(X) # 모델에다가 input 넣기.
            loss = criterion(predictions,y) # loss연산.
            running_loss += loss.item() # loss을 누적.
            _, predicted = torch.max(predictions.data, 1) # 가장 큰 값의 인덱스를 예측값으로 사용.
            correct += (predicted == y).sum().item() # 정답을 카운트.
            total += y.size(0) # 전체 개수를 카운트.
    return running_loss / len(test_loader), 100 * correct / total # 평균 loss과 정확도를 반횐.

def main():
    # Main에서는 model, optimizer, Lossfunc를 가져오기
    model = WithoutNNMLP(4, 100, 3) #이전에 정의한 MLP모델 인스턴스로 가져오기
    criterion = WithoutNNCrossEntropyLoss() #CrossEntropy를 인스턴스로 가져오기
    optimizer = torch.optim.SGD(model.parameters(), lr=0.01) #SGD를 torch에서 optimizer로 가져오기

    train_loss_list, test_loss_list, test_acc_list = [], [], []
    for epoch in range(50):
        train_loss = train(model, criterion, optimizer, train_loader)
        train_loss_list.append(train_loss)
        test_loss, test_acc = test(model, test_loader, criterion)
        test_loss_list.append(test_loss)
        test_acc_list.append(test_acc)
        print(
            f"Epoch: {epoch + 1}, Train Loss: {train_loss:.4f}, Test Loss: {test_loss:.4f}, Test Acc: {test_acc:.4f}"
        )

    plt.plot(train_loss_list, label="train loss")
    plt.plot(test_loss_list, label="test loss")
    plt.legend()
    plt.show()

    plt.plot(test_acc_list, label="test acc")
    plt.legend()
    plt.show()

    return train_loss_list, test_loss_list, test_acc_list

train_loss_list, test_loss_list, test_acc_list = main()

import torch.nn as nn
# input_size = 1, output_size = 3인 Linear_layer를 정의
nn_linear = nn.Linear(1,3)
#relu정의
nn_relu = nn.ReLU()

import torch.nn as nn

class MLP(nn.Module):
    def __init__(self, in_features, hidden_features, out_features):
        # lienar layer와 relu 정의해주기
        super().__init__()   # 상속받은 부모 클래스의 생성자 호출
        self.linear1 = nn.Linear(in_features, hidden_features) 
        self.relu = nn.ReLU() 
        self.linear2 = nn.Linear(hidden_features, out_features) 

    def forward(self, x):
        # Linear1, relu, linear2 순서대로 쌓아주기
        y = self.linear1(x)
        z = self.relu(y)
        return self.linear2(z)

# nn.Module을 이용해서 train하는 함수
def train_nn(model, criterion, optimizer, train_loader):
    model.train()  # 모델을 학습 모드로 설정
    running_loss = 0
    for X, y in train_loader:  
        optimizer.zero_grad() # 그래디언트를 0으로 초기화
        outputs = model(X) # 모델의 출력을 계산
        loss = criterion(outputs, y) # 손실을 계산
        loss.backward() # 역전파를 수행
        optimizer.step() # 가중치를 업데이트
        running_loss += loss.item() # 손실을 누적
    return running_loss / len(train_loader)  # 평균 손실을 반환

# nn.Module을 이용해서 test하는 함수
def test_nn(model, test_loader, criterion):
    model.eval()  # 모델을 평가 모드로 설정
    running_loss = 0
    correct = 0
    total = 0
    with torch.no_grad():  # 그래디언트 계산을 비활성화
        for X, y in test_loader:  # 테스트 데이터를 반복
            predictions = model(X) # 모델의 출력을 계산
            loss = criterion(predictions, y) # 손실을 계산
            running_loss += loss.item() # 손실을 누적
            _, predicted = torch.max(predictions,1)  # 가장 높은 확률을 가진 클래스를 예측
            correct += (predicted == y).sum().item() # 정확한 예측의 수를 누적
            total += y.size(0) # 전체 레이블의 수를 누적
    return running_loss / len(test_loader), 100 * correct / total  # 평균 손실과 정확도를 반환

# 모델을 돌리는 main함수
def main_nn():
    model = MLP(4, 100, 3)  # 모델을 생성
    criterion = nn.CrossEntropyLoss()  # 손실 함수를 정의
    optimizer = torch.optim.SGD(model.parameters(), lr=0.01)  # 최적화 알고리즘을 정의

    train_loss_list, test_loss_list, test_acc_list = (
        [],
        [],
        [],
    )  # 손실과 정확도를 저장할 리스트
    for epoch in range(50):  # 50 에포크 동안 학습
        train_loss = train_nn(
            model, criterion, optimizer, train_loader
        )  # 학습 손실을 계산
        train_loss_list.append(train_loss)  # 학습 손실을 저장
        test_loss, test_acc = test_nn(
            model, test_loader, criterion
        )  # 테스트 손실과 정확도를 계산
        test_loss_list.append(test_loss)  # 테스트 손실을 저장
        test_acc_list.append(test_acc)  # 테스트 정확도를 저장
        print(
            f"Epoch: {epoch + 1}, Train Loss: {train_loss:.4f}, Test Loss: {test_loss:.4f}, Test Acc: {test_acc:.4f}"
        )  # 에포크, 학습 손실, 테스트 손실, 테스트 정확도를 출력

    # 학습 손실과 테스트 손실을 그래프로 그림
    plt.plot(train_loss_list, label="train loss")
    plt.plot(test_loss_list, label="test loss")
    plt.legend()
    plt.show()

    # 테스트 정확도를 그래프로 그림
    plt.plot(test_acc_list, label="test acc")
    plt.legend()
    plt.show()

    return train_loss_list, test_loss_list, test_acc_list

weights = torch.randn(input_dim, output_dim)
bias = torch.zeros(output_dim)