import os
import csv
import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np
import pandas as pd
from sklearn.preprocessing import StandardScaler
from sklearn.model_selection import train_test_split
from sklearn.ensemble import RandomForestRegressor
from sklearn.metrics import mean_squared_error, mean_absolute_error, r2_score
from datetime import datetime, timedelta
import xgboost as xgb
import random
import matplotlib.pyplot as plt

# �쒕뜡 �쒕뱶 �ㅼ젙
random.seed(42)
np.random.seed(42)
torch.manual_seed(42)
if torch.cuda.is_available():
    torch.cuda.manual_seed_all(42)

# 蹂듦텒 踰덊샇 �낅젰 �⑥닔
def input_lotto_numbers():
    print("蹂듦텒 踰덊샇瑜� 6媛� �낅젰�섏꽭�� (�쇳몴濡� 援щ텇): ")
    user_input = input()
    numbers = [int(num.strip()) for num in user_input.split(',')]
    return numbers

# 蹂듦텒 踰덊샇 ���� �⑥닔
def save_lotto_numbers(file_path, numbers):
    date, round_number = generate_new_round(file_path)

    # �덈줈�� �곗씠�� �� �앹꽦
    new_row = [date, round_number] + numbers
    
    # 湲곗〈 �곗씠�� �쎄린
    if os.path.exists(file_path):
        with open(file_path, 'r', encoding='utf-8-sig') as file:
            reader = csv.reader(file)
            rows = list(reader)
    else:
        # �뚯씪�� �놁쑝硫� �ㅻ뜑 �앹꽦
        rows = [['異붿꺼��', '�뚯감', '踰덊샇1', '踰덊샇2', '踰덊샇3', '踰덊샇4', '踰덊샇5', '踰덊샇6']]
    
    # �� �곗씠�곕� 留� �욎뿉 �쎌엯
    rows.insert(1, new_row)

    # �뚯씪�� �ㅼ떆 �곌린
    with open(file_path, 'w', newline='', encoding='utf-8-sig') as file:
        writer = csv.writer(file)
        writer.writerows(rows)
    
    print(f"{round_number}�뚯감�� 蹂듦텒 踰덊샇媛� {file_path}�� ���λ릺�덉뒿�덈떎.")

# �뚯감�� �좎쭨 �앹꽦 �⑥닔
def generate_new_round(file_path):
    if os.path.exists(file_path) and os.path.getsize(file_path) > 0:
        with open(file_path, 'r', encoding='utf-8-sig') as file:
            reader = csv.reader(file)
            next(reader)  # �ㅻ뜑 �ㅽ궢
            last_row = next(reader)  # 泥� �곗씠�� �됱쓣 媛��몄샃�덈떎 (媛��� 理쒓렐 �뚯감)
            
            last_round = int(last_row[1])
            last_date = datetime.strptime(last_row[0], "%Y-%m-%d")
            new_round = last_round + 1
            new_date = last_date + timedelta(days=7)  # 留ㅼ＜ 蹂듦텒
            return new_date.strftime("%Y-%m-%d"), new_round

    # �뚯씪�� �녾굅�� 鍮꾩뼱�덉쑝硫� �꾩옱 �좎쭨�� 泥� �뚯감濡� �쒖옉
    return datetime.now().strftime("%Y-%m-%d"), 1

# �좉꼍留� 紐⑤뜽 �뺤쓽 (�� 源딄퀬 �볦� �ㅽ듃�뚰겕)
class ImprovedLottoModel(nn.Module):
    def __init__(self):
        super(ImprovedLottoModel, self).__init__()
        self.fc1 = nn.Linear(16, 1024)  # �� �� �덉씠�� �ш린
        self.fc2 = nn.Linear(1024, 512)
        self.fc3 = nn.Linear(512, 256)
        self.fc4 = nn.Linear(256, 128)
        self.fc5 = nn.Linear(128, 64)
        self.fc6 = nn.Linear(64, 32)
        self.fc7 = nn.Linear(32, 6)
        self.dropout = nn.Dropout(p=0.5)  # �쒕∼�꾩썐 鍮꾩쑉 議곗젙

    def forward(self, x):
        x = torch.relu(self.fc1(x))
        x = self.dropout(x)
        x = torch.relu(self.fc2(x))
        x = self.dropout(x)
        x = torch.relu(self.fc3(x))
        x = self.dropout(x)
        x = torch.relu(self.fc4(x))
        x = torch.relu(self.fc5(x))
        x = torch.relu(self.fc6(x))
        x = self.fc7(x)
        return x

# �ㅼ쨷 紐⑤뜽 �숈긽釉� �대옒�� �뺤쓽
class EnsembleModel:
    def __init__(self, models):
        self.models = models

    def predict(self, X):
        predictions = []
        for model in self.models:
            predictions.append(model.predict(X))
        return np.mean(predictions, axis=0)

# �곗씠�� 以�鍮� �⑥닔 (�쇱쿂 �붿��덉뼱留� �ы븿)
def prepare_data_with_previous_round(df):
    df = df[pd.to_numeric(df['�뚯감'], errors='coerce').notnull()]
    df['�뚯감'] = df['�뚯감'].astype(int)
    
    X, y = [], []

    for i in range(1, len(df)):
        previous_numbers = df.iloc[i-1][['踰덊샇1', '踰덊샇2', '踰덊샇3', '踰덊샇4', '踰덊샇5', '踰덊샇6']].values
        current_numbers = df.iloc[i][['踰덊샇1', '踰덊샇2', '踰덊샇3', '踰덊샇4', '踰덊샇5', '踰덊샇6']].values
        
        avg_prev = np.mean(previous_numbers)
        std_prev = np.std(previous_numbers)
        even_odd_ratio_prev = sum(1 for n in previous_numbers if n % 2 == 0) / 6.0
        sum_prev = np.sum(previous_numbers)
        product_prev = np.prod(previous_numbers)
        pairwise_diff = np.diff(previous_numbers)

        features = np.concatenate((previous_numbers, [avg_prev, std_prev, even_odd_ratio_prev, sum_prev, product_prev], pairwise_diff))
        X.append(features)
        y.append(current_numbers)
    
    X = np.array(X, dtype=np.float32)
    y = np.array(y, dtype=np.float32)

    return X, y

# �곗씠�� 利앷컯 �⑥닔 媛쒖꽑 (�� �ㅼ뼇�� �⑦꽩 異붽�)
def augment_data(X, y):
    synthetic_data_X = []
    synthetic_data_y = []
    
    for i in range(len(X)):
        for _ in range(5):  # 媛� �먮낯 �섑뵆�� ���� �ㅼ꽢 媛�吏� �⑹꽦 �섑뵆 �앹꽦
            noise = np.random.normal(0, 0.02, X.shape[1])  # �� �묒� �몄씠利� 異붽�
            synthetic_sample_X = X[i] + noise
            synthetic_data_X.append(synthetic_sample_X)
            synthetic_data_y.append(y[i])  # y 媛믩룄 �숈씪�섍쾶 利앷컯

    synthetic_X = np.array(synthetic_data_X, dtype=np.float32)
    synthetic_y = np.array(synthetic_data_y, dtype=np.float32)

    X_augmented = np.concatenate((X, synthetic_X))
    y_augmented = np.concatenate((y, synthetic_y))
    
    return X_augmented, y_augmented

# �ㅼ뼇�� 紐⑤뜽�� �ъ슜�� �숈뒿 �⑥닔��
def train_random_forest(X_train, y_train):
    rf_model = RandomForestRegressor(n_estimators=300, max_depth=50, min_samples_split=10, random_state=42)
    rf_model.fit(X_train, y_train)
    return rf_model

def train_xgboost(X_train, y_train):
    xgb_model = xgb.XGBRegressor(objective='reg:squarederror', n_estimators=200, max_depth=5, learning_rate=0.05, random_state=42)
    xgb_model.fit(X_train, y_train)
    return xgb_model

# �좉꼍留� 紐⑤뜽 �숈뒿 �⑥닔 (�숈뒿瑜� 議곗젙)
def train_model(model, X_train, y_train, X_val, y_val, max_epochs=10000, patience=500):
    criterion = nn.MSELoss()
    optimizer = optim.AdamW(model.parameters(), lr=0.0001)  # �� ��� �숈뒿瑜좉낵 AdamW �ъ슜
    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=200, eta_min=1e-6)

    best_loss = float('inf')
    training_losses = []
    validation_losses = []
    early_stop_count = 0

    save_dir = 'saved_models'
    os.makedirs(save_dir, exist_ok=True)

    for epoch in range(max_epochs):
        model.train()
        optimizer.zero_grad()
        outputs = model(X_train)
        loss = criterion(outputs, y_train)
        loss.backward()
        optimizer.step()

        training_losses.append(loss.item())

        model.eval()
        with torch.no_grad():
            val_outputs = model(X_val)
            val_loss = criterion(val_outputs, y_val)
            validation_losses.append(val_loss.item())

            if val_loss < best_loss:
                best_loss = val_loss
                model_path = os.path.join(save_dir, f'best_{model.__class__.__name__}_loss_{best_loss:.4f}.pth')
                torch.save(model.state_dict(), model_path)
                early_stop_count = 0  # 議곌린 醫낅즺 移댁슫�� 珥덇린��
            else:
                early_stop_count += 1

        scheduler.step()

        if early_stop_count >= patience:
            print(f"Early stopping at epoch {epoch}")
            break

    best_model_path = os.path.join(save_dir, f'best_{model.__class__.__name__}_loss_{best_loss:.4f}.pth')
    model.load_state_dict(torch.load(best_model_path))
    return model, training_losses, validation_losses

# 紐⑤뜽 �됯� 諛� 濡쒓퉭 �⑥닔
def evaluate_model(model, X_test, y_test):
    model.eval()
    with torch.no_grad():
        test_outputs = model(X_test)
        test_loss = nn.MSELoss()(test_outputs, y_test).item()
        test_mae = mean_absolute_error(y_test.cpu().numpy(), test_outputs.cpu().numpy())
        test_r2 = r2_score(y_test.cpu().numpy(), test_outputs.cpu().numpy())
        print(f"Test Loss (MSE): {test_loss:.4f}")
        print(f"Test MAE: {test_mae:.4f}")
        print(f"Test R짼: {test_r2:.4f}")
    return test_loss, test_mae, test_r2

# �숈뒿 諛� 寃�利� �먯떎 怨≪꽑 �뚮줈�� �⑥닔
def plot_loss_curves(training_losses, validation_losses, model_name):
    plt.figure(figsize=(10, 5))
    plt.plot(training_losses, label='Training Loss')
    plt.plot(validation_losses, label='Validation Loss')
    plt.xlabel('Epochs')
    plt.ylabel('Loss')
    plt.title(f'Loss Curves for {model_name}')
    plt.legend()
    plt.savefig(f'loss_curves_{model_name}.png')
    plt.close()

# �숈긽釉� �덉륫 �⑥닔 媛쒖꽑 (媛�以묒튂 異붽�)
def ensemble_predict_advanced(models, X_input, n_sets=10):
    predictions = []

    model_weights = [0.5, 0.3, 0.2]  # 紐⑤뜽�� �깅뒫�� �곕씪 媛�以묒튂瑜� 遺���

    for _ in range(n_sets):
        noisy_input = X_input + torch.normal(0, 0.05, size=X_input.shape).to(X_input.device)

        stacked_predictions = np.zeros((noisy_input.shape[0], 6))
        for idx, model in enumerate(models):
            if isinstance(model, nn.Module):
                prediction = model(noisy_input).detach().cpu().numpy()
            else:
                prediction = model.predict(noisy_input.cpu().numpy())
            stacked_predictions += model_weights[idx] * prediction

        final_prediction = stacked_predictions / np.sum(model_weights)

        # �덉륫 寃곌낵瑜� �ㅻ쫫李⑥닚�쇰줈 �뺣젹�섍퀬 �쒕뜡 �욊린 異붽�
        final_prediction_sorted = np.sort(np.clip(np.round(final_prediction), 1, 45).astype(int))
        np.random.shuffle(final_prediction_sorted)
        predictions.append(final_prediction_sorted)

    return predictions

# 硫붿씤 �⑥닔
def main():
    file_path = r'C:\Users\user\projects\st\lotto\lotto_numbers.csv'
    numbers = input_lotto_numbers()  # 蹂듦텒 踰덊샇瑜� �낅젰諛쏆뒿�덈떎.
    save_lotto_numbers(file_path, numbers)

    df = pd.read_csv(file_path)
    X, y = prepare_data_with_previous_round(df)

    X, y = augment_data(X, y)

    X_train, X_temp, y_train, y_temp = train_test_split(X, y, test_size=0.3, random_state=42)
    X_val, X_test, y_val, y_test = train_test_split(X_temp, y_temp, test_size=0.2, random_state=42)

    scaler = StandardScaler()
    X_train = scaler.fit_transform(X_train)
    X_val = scaler.transform(X_val)
    X_test = scaler.transform(X_test)

    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    X_train = torch.tensor(X_train, dtype=torch.float32).to(device)
    y_train = torch.tensor(y_train, dtype=torch.float32).to(device)
    X_val = torch.tensor(X_val, dtype=torch.float32).to(device)
    y_val = torch.tensor(y_val, dtype=torch.float32).to(device)
    X_test = torch.tensor(X_test, dtype=torch.float32).to(device)
    y_test = torch.tensor(y_test, dtype=torch.float32).to(device)

    improved_nn = ImprovedLottoModel().to(device)
    improved_nn, train_losses, val_losses = train_model(improved_nn, X_train, y_train, X_val, y_val)

    plot_loss_curves(train_losses, val_losses, "ImprovedLottoModel")

    rf_model = train_random_forest(X_train.cpu().numpy(), y_train.cpu().numpy())
    xgb_model = train_xgboost(X_train.cpu().numpy(), y_train.cpu().numpy())

    ensemble_model = EnsembleModel([improved_nn, rf_model, xgb_model])

    test_loss, test_mae, test_r2 = evaluate_model(improved_nn, X_test, y_test)

    X_input = X_test[0].unsqueeze(0).to(device)
    ensemble_predictions = ensemble_predict_advanced([improved_nn, rf_model, xgb_model], X_input)

    for idx, prediction in enumerate(ensemble_predictions):
        print(f"�덉륫 �명듃 {idx + 1}: {prediction}")

if __name__ == "__main__":
    main()