import torch.multiprocessing
torch.multiprocessing.set_sharing_strategy('file_system')
import warnings
warnings.filterwarnings("ignore")

import pytorch_lightning as pl
from pytorch_lightning.callbacks import ModelCheckpoint
from tools.cfg import py2cfg
import os
import torch
from torch import nn
import cv2
import numpy as np
import argparse
from pathlib import Path
from tools.metric import Evaluator
from pytorch_lightning.loggers import CSVLogger
import random

def seed_everything(seed):
    random.seed(seed)
    os.environ['PYTHONHASHSEED'] = str(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False


def get_args():
    parser = argparse.ArgumentParser()
    arg = parser.add_argument
    arg("-c", "--config_path", type=Path, help="Path to the config.", required=True)
    return parser.parse_args()


class Supervision_Train(pl.LightningModule):
    def __init__(self, config):
        super().__init__()
        self.config = config
        self.net = config.net
        self.automatic_optimization = False

        self.loss = config.loss

        self.metrics_train = Evaluator(num_class=config.num_classes)
        self.metrics_val = Evaluator(num_class=config.num_classes)

    def forward(self, x):
        # only net is used in the prediction/inference
        seg_pre = self.net(x)
        return seg_pre

    def training_step(self, batch, batch_idx):
        img, mask = batch['img'], batch['gt_semantic_seg']

        prediction = self.net(img)
        # print("instep:",prediction.shape,mask.shape,mask.device)
        # instep: torch.Size([4, 2, 512, 512]) torch.Size([4, 512, 512]) cuda:0
        # instep: torch.Size([4, 2, 512, 512]) torch.Size([4, 512, 512]) cuda:1

        return {"pred": prediction,"mask":mask}      

    def training_step_end(self,training_step_outputs):
        
        # 已经自动将分布在两块卡上的输出合并起来了
        predictions = training_step_outputs["pred"]
        masks = training_step_outputs["mask"]

        # print("step end:",predictions.shape,masks.shape,masks.device)
        # step end: torch.Size([8, 2, 512, 512]) torch.Size([8, 512, 512]) cuda:0
    
        loss = self.loss(predictions, masks)

        if self.config.use_aux_loss:
            pre_mask = nn.Softmax(dim=1)(predictions[0])
        else:
            pre_mask = nn.Softmax(dim=1)(predictions)

        # pre_mask = predictions.argmax(dim=1)#我自己的
        if isinstance(predictions, tuple):
            predictions = predictions[0]  # UNetFormer取主输出
        pre_mask = predictions.argmax(dim=1)

        for i in range(masks.shape[0]):
            self.metrics_train.add_batch(masks[i].cpu().numpy(), pre_mask[i].cpu().numpy())

        # supervision stage
        opt = self.optimizers(use_pl_optimizer=False)
        self.manual_backward(loss)
        #if (batch_idx + 1) % self.config.accumulate_n == 0:
        opt.step()
        opt.zero_grad()

        sch = self.lr_schedulers()
        if self.trainer.is_last_batch and (self.trainer.current_epoch + 1) % 1 == 0:
            sch.step()

        return {"loss": loss}


    def training_epoch_end(self, outputs):
        if 'vaihingen' in self.config.log_name:
            mIoU = np.nanmean(self.metrics_train.Intersection_over_Union()[:-1])
            F1 = np.nanmean(self.metrics_train.F1()[:-1])
        elif 'potsdam' in self.config.log_name:
            mIoU = np.nanmean(self.metrics_train.Intersection_over_Union()[:-1])
            F1 = np.nanmean(self.metrics_train.F1()[:-1])
        elif 'whu' in self.config.log_name:
            mIoU = np.nanmean(self.metrics_train.Intersection_over_Union()[:-1])
            F1 = np.nanmean(self.metrics_train.F1()[:-1])
        elif 'mass' in self.config.log_name:
            mIoU = np.nanmean(self.metrics_train.Intersection_over_Union()[:-1])
            F1 = np.nanmean(self.metrics_train.F1()[:-1])
        elif 'inria' in self.config.log_name:
            mIoU = np.nanmean(self.metrics_train.Intersection_over_Union()[:-1])
            F1 = np.nanmean(self.metrics_train.F1()[:-1])
        else:
            mIoU = np.nanmean(self.metrics_train.Intersection_over_Union())
            F1 = np.nanmean(self.metrics_train.F1())

        OA = np.nanmean(self.metrics_train.OA())
        iou_per_class = self.metrics_train.Intersection_over_Union()
        eval_value = {'mIoU':np.round(mIoU,6),'F1': np.round(F1,6),'OA': np.round(OA,6)}
        print(" ")
        print('train:', eval_value)

        iou_value = {}
        for class_name, iou in zip(self.config.classes, iou_per_class):
            iou_value[class_name] = np.round(iou,6)
        #print(iou_value)
        print("======================")
        print(" ")

        self.metrics_train.reset()
        loss = torch.stack([x["loss"] for x in outputs]).mean()
        log_dict = {"train_loss": loss, 'train_mIoU': mIoU, 'train_F1': F1, 'train_OA': OA}
        self.log_dict(log_dict, prog_bar=True)

    def validation_step(self, batch, batch_idx):
        img, mask = batch['img'], batch['gt_semantic_seg']
        prediction = self.forward(img)

        # print(prediction.shape,prediction.device)
        # torch.Size([4, 2, 512, 512]) cuda:1
        # torch.Size([4, 2, 512, 512]) cuda:0

        return {"pred": prediction,"mask":mask}   

    def validation_step_end(self,validation_step_outputs):

        predictions = validation_step_outputs["pred"]
        masks = validation_step_outputs["mask"]

        # print("valid step end:",predictions.shape,predictions.device)
        # valid step end: torch.Size([8, 2, 512, 512]) cuda:0

        pre_mask = nn.Softmax(dim=1)(predictions)
        pre_mask = pre_mask.argmax(dim=1)
        for i in range(masks.shape[0]):
            self.metrics_val.add_batch(masks[i].cpu().numpy(), pre_mask[i].cpu().numpy())

        loss_val = self.loss(predictions, masks)
        return {"loss_val": loss_val}

    def validation_epoch_end(self, outputs):
        if 'vaihingen' in self.config.log_name:
            mIoU = np.nanmean(self.metrics_val.Intersection_over_Union()[:-1])
            F1 = np.nanmean(self.metrics_val.F1()[:-1])
        elif 'potsdam' in self.config.log_name:
            mIoU = np.nanmean(self.metrics_val.Intersection_over_Union()[:-1])
            F1 = np.nanmean(self.metrics_val.F1()[:-1])
        elif 'whu' in self.config.log_name:
            mIoU = np.nanmean(self.metrics_val.Intersection_over_Union()[:-1])
            F1 = np.nanmean(self.metrics_val.F1()[:-1])
        elif 'mass' in self.config.log_name:
            mIoU = np.nanmean(self.metrics_val.Intersection_over_Union()[:-1])
            F1 = np.nanmean(self.metrics_val.F1()[:-1])
        elif 'inria' in self.config.log_name:
            mIoU = np.nanmean(self.metrics_val.Intersection_over_Union()[:-1])
            F1 = np.nanmean(self.metrics_val.F1()[:-1])
        else:
            mIoU = np.nanmean(self.metrics_val.Intersection_over_Union())
            F1 = np.nanmean(self.metrics_val.F1())

        OA = np.nanmean(self.metrics_val.OA())
        iou_per_class = self.metrics_val.Intersection_over_Union()

        eval_value = {'mIoU':np.around(mIoU,6) ,
                      'F1': np.around(F1,6),
                      'OA': np.around(OA,6)}

        print(" ")
        print('val:', eval_value)

        iou_value = {}
        for class_name, iou in zip(self.config.classes, iou_per_class):
            iou_value[class_name] = np.around(iou,6)
        print(iou_value)
        print("======================")
        
        self.metrics_val.reset()
        loss = torch.stack([x["loss_val"] for x in outputs]).mean()
        log_dict = {"val_loss": loss, 'val_mIoU': mIoU, 'val_F1': F1, 'val_OA': OA}
        self.log_dict(log_dict, prog_bar=True)
        # Also save a standalone model weight file after validation
        # This ensures we have model weights even if checkpoint saving fails
        if hasattr(self.config, 'weights_path'):
            model_save_path = os.path.join(self.config.weights_path, f"epoch_{self.current_epoch}_model_F1_{F1:.4f}_mIoU_{mIoU:.4f}.pth")
            torch.save(self.net.state_dict(), model_save_path)
            print(f"Saved model weights to {model_save_path}")

    def configure_optimizers(self):
        optimizer = self.config.optimizer
        lr_scheduler = self.config.lr_scheduler

        return [optimizer], [lr_scheduler]

    def train_dataloader(self):

        return self.config.train_loader

    def val_dataloader(self):

        return self.config.val_loader


# # training
# def main():
#     args = get_args()
#     config = py2cfg(args.config_path)
#     seed_everything(42)
#
#     checkpoint_callback = ModelCheckpoint(save_top_k=config.save_top_k, monitor=config.monitor,
#                                           save_last=config.save_last, mode=config.monitor_mode,
#                                           dirpath=config.weights_path,
#                                           filename=config.weights_name)
#     logger = CSVLogger('lightning_logs', name=config.log_name)
#
#     model = Supervision_Train(config)
#     if config.pretrained_ckpt_path:
#         model = Supervision_Train.load_from_checkpoint(config.pretrained_ckpt_path, config=config)
#
#     trainer = pl.Trainer(devices=config.gpus, max_epochs=config.max_epoch, accelerator='gpu',
#                          check_val_every_n_epoch=config.check_val_every_n_epoch,
#                          callbacks=[checkpoint_callback], strategy='auto',
#                          resume_from_checkpoint=config.resume_ckpt_path, logger=logger,
#                          precision=32
#                          )
#
#     trainer.fit(model=model)
# training
def main():
    args = get_args()
    config = py2cfg(args.config_path)
    seed_everything(42)

    # Ensure weights directory exists
    os.makedirs(config.weights_path, exist_ok=True)

    # Create more verbose filename pattern with metrics
    filename_pattern = "{epoch:02d}-{val_F1:.4f}-{val_mIoU:.4f}"

    # Setup checkpoint callback with explicit filename
    checkpoint_callback = ModelCheckpoint(
        dirpath=config.weights_path,
        filename=filename_pattern,
        save_top_k=config.save_top_k,
        verbose=True,
        monitor=config.monitor,
        mode=config.monitor_mode,
        save_last=config.save_last
    )

    # Make sure log directory exists
    log_dir = os.path.dirname(config.log_name)
    os.makedirs(log_dir, exist_ok=True)

    # Create the logger with appropriate path
    logger = CSVLogger(
        save_dir=os.path.dirname(config.log_name),
        name=os.path.basename(config.log_name),
        version=None  # Let Lightning handle versioning
    )

    model = Supervision_Train(config)
    if config.pretrained_ckpt_path:
        model = Supervision_Train.load_from_checkpoint(config.pretrained_ckpt_path, config=config)

    # Add additional useful callbacks if needed
    callbacks = [checkpoint_callback]

    trainer = pl.Trainer(
        devices=config.gpus,
        max_epochs=config.max_epoch,
        accelerator='gpu',
        check_val_every_n_epoch=config.check_val_every_n_epoch,
        callbacks=callbacks,
        strategy='auto',
        resume_from_checkpoint=config.resume_ckpt_path,
        logger=logger,
        precision=32,
        log_every_n_steps=10  # More frequent logging
    )

    # Print confirmation message
    print(f"Model checkpoints will be saved to: {config.weights_path}")
    print(f"Using checkpoint filename pattern: {filename_pattern}")

    trainer.fit(model=model)

    # Print final confirmation
    print(f"Training completed. Check {config.weights_path} for saved checkpoints.")

    # Save final model weights as a standalone file as well
    final_model_path = os.path.join(config.weights_path, "final_model.pth")
    torch.save(model.net.state_dict(), final_model_path)
    print(f"Saved final model weights to {final_model_path}")


if __name__ == "__main__":
    main()

