import os
import sys
import torch
import torch.nn as nn
from datetime import datetime
from torch.nn.utils import clip_grad_norm_
import torch.optim as optim
from torch.utils.data import DataLoader
import json

# إعداد المسارات
dir_path = os.path.dirname(os.path.realpath(__file__))
sys.path.append(os.path.join(dir_path, '..'))

from nets.model import Model
from Actor.actor import Actor
from dataloader import VRP_Dataset
from google_solver.google_model import evaluate_google_model

# تحميل الإعدادات
with open('params.json', 'r') as f:
    params = json.load(f)

# تحديد مجلد آمن لحفظ الملفات داخل $HOME
safe_data_dir = "/home/user/data"
os.makedirs(safe_data_dir, exist_ok=True)

# حفظ نسخة من الإعدادات
with open(os.path.join(safe_data_dir, 'params_saved.json'), 'w') as f:
    json.dump(params, f)

# إعداد المتغيرات
device = params['device']
run_tests = params['run_tests']
save_results = params['save_results']
dataset_path = params['dataset_path']

# حجم البيانات
train_dataset_size = params['train_dataset_size']
validation_dataset_size = params['validation_dataset_size']
baseline_dataset_size = params['baseline_dataset_size']

# خصائص المشكلة
num_nodes = params['num_nodes']
num_depots = params['num_depots']
embedding_size = params['embedding_size']
sample_size = params['sample_size']
gradient_clipping = params['gradient_clipping']
num_neighbors_encoder = params['num_neighbors_encoder']
num_neighbors_action = params['num_neighbors_action']
num_movers = params['num_movers']
learning_rate = params['learning_rate']
batch_size = params['batch_size']
test_batch_size = params['test_batch_size']

# التحميل
validation_dataset = VRP_Dataset(validation_dataset_size, num_nodes, num_depots, dataset_path, device)
baseline_dataset = VRP_Dataset(train_dataset_size, num_nodes, num_depots, dataset_path, device)
if params['overfit_test']:
    train_dataset = baseline_dataset = validation_dataset = VRP_Dataset(train_dataset_size, num_nodes, num_depots, dataset_path, device)

# تقييم Google OR-Tools
google_scores = evaluate_google_model(validation_dataset)
tot_google_scores = google_scores.sum().item()
input_size = validation_dataset.model_input_length()

# تعريف النماذج
model = Model(input_size=input_size, embedding_size=embedding_size, decoder_input_size=params["decoder_input_size"])
actor = Actor(model=model, num_movers=num_movers, num_neighbors_encoder=num_neighbors_encoder,
              num_neighbors_action=num_neighbors_action, device=device, normalize=False)
actor.train_mode()

baseline_model = Model(input_size=input_size, embedding_size=embedding_size, decoder_input_size=params["decoder_input_size"])
baseline_actor = Actor(model=baseline_model, num_movers=num_movers, num_neighbors_encoder=num_neighbors_encoder,
                       num_neighbors_action=num_neighbors_action, device=device, normalize=False)
baseline_actor.greedy_search()
baseline_actor.load_state_dict(actor.state_dict())

nn_actor = Actor(model=None, num_movers=1, num_neighbors_action=1, device=device)
nn_actor.nearest_neighbors()

optimizer = optim.Adam(params=actor.parameters(), lr=learning_rate)

# ملفات الإخراج
train_results_file = os.path.join(safe_data_dir, "train_results.txt")
test_results_file = os.path.join(safe_data_dir, "test_results.txt")
model_path = os.path.join(safe_data_dir, "model_state_dict.pt")
optimizer_path = os.path.join(safe_data_dir, "optimizer_state_dict.pt")

train_batch_record = 100
validation_record = 100
baseline_record = None

# التدريب
for epoch in range(params['num_epochs']):
    if not params['overfit_test']:
        train_dataset = VRP_Dataset(train_dataset_size, num_nodes, num_depots, dataset_path, device)

    train_dataloader = DataLoader(train_dataset, batch_size=batch_size, collate_fn=train_dataset.collate)
    for i, batch in enumerate(train_dataloader):
        with torch.no_grad():
            nn_output = nn_actor(batch)
            tot_nn_cost = nn_output['total_time'].sum().item()

            baseline_actor.greedy_search()
            baseline_cost = baseline_actor(batch)['total_time']

        actor.train_mode()
        actor_output = actor(batch)
        actor_cost, log_probs = actor_output['total_time'], actor_output['log_probs']

        loss = ((actor_cost - baseline_cost).detach() * log_probs).mean()

        optimizer.zero_grad()
        loss.backward()
        if gradient_clipping:
            for group in optimizer.param_groups:
                clip_grad_norm_(group['params'], 1, norm_type=2)
        optimizer.step()

        tot_actor_cost = actor_cost.sum().item()
        tot_baseline_cost = baseline_cost.sum().item()
        actor_nn_ratio = tot_actor_cost / tot_nn_cost
        actor_baseline_ratio = tot_actor_cost / tot_baseline_cost
        train_batch_record = min(train_batch_record, actor_nn_ratio)

        result = f"{epoch}, {i}, {actor_nn_ratio:.4f}, {actor_baseline_ratio:.4f}, {train_batch_record:.4f}"
        print(result)
        if save_results:
            with open(train_results_file, 'a') as f:
                f.write(result + '\n')
        del batch

    # التحقق من الأداء
    if epoch % 5 == 0:
        baseline_dataloader = DataLoader(baseline_dataset, batch_size=batch_size, collate_fn=baseline_dataset.collate)
        tot_cost = []
        for batch in baseline_dataloader:
            with torch.no_grad():
                actor.greedy_search()
                cost = actor(batch)['total_time']
            tot_cost.append(cost)
        tot_cost = torch.cat(tot_cost, dim=0)
        if baseline_record is None or (tot_cost < baseline_record).float().mean().item() > 0.9:
            baseline_record = tot_cost
            baseline_actor.load_state_dict(actor.state_dict())
            print('\nNew baseline record\n')

    # التقييم وحفظ النموذج دائمًا
    if (epoch % 10 == 0) and run_tests:
        b = max(int(batch_size // sample_size**2), 1)
        validation_dataloader = DataLoader(validation_dataset, batch_size=b, collate_fn=validation_dataset.collate)

        tot_cost = 0
        tot_nn_cost = 0
        for batch in validation_dataloader:
            with torch.no_grad():
                actor.beam_search(sample_size)
                actor_output = actor(batch)
                cost = actor_output['total_time']

                nn_output = nn_actor(batch)
                nn_cost = nn_output['total_time']

            tot_cost += cost.sum().item()
            tot_nn_cost += nn_cost.sum().item()

        ratio = tot_cost / tot_nn_cost
        validation_record = min(validation_record, ratio)
        actor_google_ratio = tot_cost / tot_google_scores
        print(f"\nTest results:\nActor/Google: {actor_google_ratio:.4f}, Actor/NN: {ratio:.4f}, Best NN Ratio: {validation_record:.4f}\n")

        if save_results:
            with open(test_results_file, 'a') as f:
                f.write(f"{epoch}, {actor_google_ratio:.4f}, {ratio:.4f}, {validation_record:.4f}\n")
            torch.save(actor.state_dict(), model_path)
            torch.save(optimizer.state_dict(), optimizer_path)

print("End.")