import ttach as tta
import multiprocessing.pool as mpp
import multiprocessing as mp
import time
from train_supervision import *
import argparse
from pathlib import Path
import cv2
import numpy as np
import torch

from torch import nn
from torch.utils.data import DataLoader
from tqdm import tqdm


def label2rgb(mask):
    h, w = mask.shape[0], mask.shape[1]
    mask_rgb = np.zeros(shape=(h, w, 3), dtype=np.uint8)
    mask_convert = mask[np.newaxis, :, :]
    mask_rgb[np.all(mask_convert == 0, axis=0)] = [255, 255, 255]
    mask_rgb[np.all(mask_convert == 1, axis=0)] = [255, 0, 0]
    mask_rgb[np.all(mask_convert == 2, axis=0)] = [255, 255, 0]
    mask_rgb[np.all(mask_convert == 3, axis=0)] = [0, 0, 255]
    mask_rgb[np.all(mask_convert == 4, axis=0)] = [159, 129, 183]
    mask_rgb[np.all(mask_convert == 5, axis=0)] = [0, 255, 0]
    mask_rgb[np.all(mask_convert == 6, axis=0)] = [255, 195, 128]
    return mask_rgb


def img_writer(inp):
    (mask,  mask_id, rgb) = inp
    if rgb:
        mask_name_tif = mask_id + '.png'
        mask_tif = label2rgb(mask)
        mask_tif = cv2.cvtColor(mask_tif, cv2.COLOR_RGB2BGR)
        cv2.imwrite(mask_name_tif, mask_tif)
    else:
        mask_png = mask.astype(np.uint8)
        mask_name_png = mask_id + '.png'
        cv2.imwrite(mask_name_png, mask_png)


def get_args():
    parser = argparse.ArgumentParser()
    arg = parser.add_argument
    arg("-c", "--config_path", type=Path, required=True, help="Path to  config")
    arg("-o", "--output_path", type=Path, help="Path where to save resulting masks.", required=True)
    arg("-t", "--tta", help="Test time augmentation.", default=None, choices=[None, "d4", "lr"]) ## lr is flip TTA, d4 is multi-scale TTA
    arg("--rgb", help="whether output rgb masks", action='store_true')
    arg("--val", help="whether eval validation set", action='store_true')
    return parser.parse_args()

#
# def main():
#     args = get_args()
#     config = py2cfg(args.config_path)
#     args.output_path.mkdir(exist_ok=True, parents=True)
#
#     # model = Supervision_Train.load_from_checkpoint(os.path.join(config.weights_path, config.test_weights_name+'.ckpt'), config=config)
#     model = Supervision_Train(config)
#     # Then load the state dict
#     state_dict = torch.load(os.path.join(config.weights_path, config.test_weights_name + '.pth'))
#     model.net.load_state_dict(state_dict)
#     model.cuda(config.gpus[0])
#
#     model.eval()
#     if args.tta == "lr":
#         transforms = tta.Compose(
#             [
#                 tta.HorizontalFlip(),
#                 tta.VerticalFlip()
#             ]
#         )
#         model = tta.SegmentationTTAWrapper(model, transforms)
#     elif args.tta == "d4":
#         transforms = tta.Compose(
#             [
#                 tta.HorizontalFlip(),
#                 tta.VerticalFlip(),
#                 #tta.Rotate90(angles=[90, 180, 270]),
#                 tta.Scale(scales=[0.75, 1.0, 1.25, 1.5], interpolation='bicubic', align_corners=False),
#                 #tta.Resize((512,512),(1024,1024),interpolation='bicubic', align_corners=False)
#                 # tta.Multiply(factors=[0.8, 1, 1.2])
#             ]
#         )
#         model = tta.SegmentationTTAWrapper(model, transforms)
#
#     test_dataset = config.test_dataset
#     if args.val:
#         evaluator = Evaluator(num_class=config.num_classes)
#         evaluator.reset()
#         test_dataset = config.val_dataset
#
#     with torch.no_grad():
#         test_loader = DataLoader(
#             test_dataset,
#             batch_size=1,
#             num_workers=4,
#             pin_memory=True,
#             drop_last=False,
#         )
#         results = []
#         for input in tqdm(test_loader):
#             # raw_prediction NxCxHxW
#             raw_predictions = model(input['img'].cuda(config.gpus[0]))
#
#             image_ids = input["img_id"]
#             if args.val:
#                 masks_true = input['gt_semantic_seg']
#
#             img_type = input['img_type']
#
#             raw_predictions = nn.Softmax(dim=1)(raw_predictions)
#             predictions = raw_predictions.argmax(dim=1)
#
#             for i in range(raw_predictions.shape[0]):
#                 mask = predictions[i].cpu().numpy()
#                 mask_name = image_ids[i]
#                 mask_type = img_type[i]
#                 if args.val:
#                     if not os.path.exists(os.path.join(args.output_path, mask_type)):
#                         os.mkdir(os.path.join(args.output_path, mask_type))
#                     evaluator.add_batch(pre_image=mask, gt_image=masks_true[i].cpu().numpy())
#                     results.append((mask, str(args.output_path / mask_type / mask_name), args.rgb))
#                 else:
#                     results.append((mask, str(args.output_path / mask_name), args.rgb))
#     if args.val:
#         iou_per_class = evaluator.Intersection_over_Union()
#         f1_per_class = evaluator.F1()
#         OA = evaluator.OA()
#         for class_name, class_iou, class_f1 in zip(config.classes, iou_per_class, f1_per_class):
#             print('F1_{}:{}, IOU_{}:{}'.format(class_name, class_f1, class_name, class_iou))
#         print('F1:{}, mIOU:{}, OA:{}'.format(np.nanmean(f1_per_class), np.nanmean(iou_per_class), OA))
#
#     t0 = time.time()
#     mpp.Pool(processes=mp.cpu_count()).map(img_writer, results)
#     t1 = time.time()
#     img_write_time = t1 - t0
#     print('images writing spends: {} s'.format(img_write_time))
#
#
# if __name__ == "__main__":
#     main()
def main() -> None:
    seed_everything(42)

    args = get_args()
    config = py2cfg(args.config_path)
    args.output_path.mkdir(exist_ok=True, parents=True)

    # --------------------------- 模型加载 ---------------------------------------
    model = Supervision_Train(config)
    state_dict_path = os.path.join(
        config.weights_path, f"{config.test_weights_name}.pth"
    )
    state_dict = torch.load(state_dict_path, map_location="cpu")
    model.net.load_state_dict(state_dict)
    model = model.cuda(config.gpus[0]).eval()

    # --------------------------- TTA 包装 ---------------------------------------
    if args.tta == "lr":
        transforms = tta.Compose([tta.HorizontalFlip(), tta.VerticalFlip()])
        model = tta.SegmentationTTAWrapper(model, transforms)
    elif args.tta == "d4":
        transforms = tta.Compose(
            [
                tta.HorizontalFlip(),
                tta.VerticalFlip(),
                tta.Scale(
                    scales=[0.75, 1.0, 1.25, 1.5],
                    interpolation="bicubic",
                    align_corners=False,
                ),
            ]
        )
        model = tta.SegmentationTTAWrapper(model, transforms)
    # ---------------------------------------------------------------------------

    # --------------------------- 数据集 & 评估器 --------------------------------
    test_dataset = config.test_dataset
    if args.val:
        evaluator = Evaluator(num_class=config.num_classes)
        evaluator.reset()
        test_dataset = config.val_dataset

    test_loader = DataLoader(
        test_dataset,
        batch_size=1,
        num_workers=4,
        pin_memory=True,
        drop_last=False,
    )

    # ---------------------------------------------------------------------------
    results = []
    total_images = 0

    # --------------------------- 计时：推理开始 ---------------------------------
    t_infer_start = time.time()

    with torch.no_grad():
        for batch in tqdm(test_loader, desc="Inference"):
            logits = model(batch["img"].cuda(config.gpus[0]))
            softmaxed = nn.Softmax(dim=1)(logits)
            preds = softmaxed.argmax(dim=1)

            image_ids = batch["img_id"]
            img_type = batch["img_type"]

            if args.val:
                masks_true = batch["gt_semantic_seg"]

            for i in range(preds.size(0)):
                mask = preds[i].cpu().numpy()
                name = image_ids[i]
                mtype = img_type[i]

                if args.val:
                    save_dir = args.output_path / mtype
                    save_dir.mkdir(exist_ok=True, parents=True)
                    evaluator.add_batch(
                        pre_image=mask, gt_image=masks_true[i].cpu().numpy()
                    )
                    results.append(
                        (mask, str(save_dir / name), args.rgb)
                    )
                else:
                    results.append(
                        (mask, str(args.output_path / name), args.rgb)
                    )

            total_images += preds.size(0)

    # --------------------------- 计时：推理结束 ---------------------------------
    t_infer_end = time.time()
    infer_time = t_infer_end - t_infer_start
    infer_fps = total_images / infer_time if infer_time > 0 else 0.0
    print(
        f"Inference time: {infer_time:.2f} s for {total_images} images "
        f"→ FPS: {infer_fps:.2f}"
    )

    # --------------------------- 评估指标 ----------------------------------------
    if args.val:
        ious = evaluator.Intersection_over_Union()
        f1s = evaluator.F1()
        oa = evaluator.OA()

        for cls, iou, f1 in zip(config.classes, ious, f1s):
            print(f"F1_{cls}: {f1:.4f}, IOU_{cls}: {iou:.4f}")
        print(f"Mean   F1 : {np.nanmean(f1s):.4f}")
        print(f"Mean  IOU : {np.nanmean(ious):.4f}")
        print(f"Overall OA: {oa:.4f}")

    # --------------------------- 写图（多进程） ----------------------------------
    t_write_start = time.time()
    with mpp.Pool(processes=mp.cpu_count()) as pool:
        pool.map(img_writer, results)
    t_write_end = time.time()
    write_time = t_write_end - t_write_start
    write_fps = total_images / write_time if write_time > 0 else 0.0
    print(
        f"Image writing time: {write_time:.2f} s "
        f"→ Write FPS: {write_fps:.2f}"
    )

    # --------------------------- 综合统计 ----------------------------------------
    overall_time = infer_time + write_time
    overall_fps = total_images / overall_time if overall_time > 0 else 0.0
    print(
        f"Total runtime (infer + write): {overall_time:.2f} s "
        f"→ Overall FPS: {overall_fps:.2f}"
    )


if __name__ == "__main__":
    main()