import os
import gradio as gr
from detectron2 import model_zoo
from detectron2.engine import DefaultTrainer
from detectron2.config import get_cfg
from detectron2.data.datasets import register_coco_instances
from detectron2.utils.visualizer import Visualizer
import cv2


# Step 1: Register the training dataset
def register_dataset(train_json, train_images):
    """
    Registers the training dataset with Detectron2.
    """
    register_coco_instances("floorplan_train", {}, train_json, train_images)
    print("Dataset registered!")

# Step 2: Configure the model
def get_training_config(output_dir, num_classes):
    """
    Returns the configuration for training Mask R-CNN.
    """
    cfg = get_cfg()
    cfg.MODEL.DEVICE = "cpu"
    cfg.merge_from_file(model_zoo.get_config_file("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml"))

    # Dataset and model configuration
    cfg.DATASETS.TRAIN = ("floorplan_train",)
    cfg.DATALOADER.NUM_WORKERS = 4
    cfg.MODEL.WEIGHTS = "detectron2://COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x/137849600/model_final_f10217.pkl"
    cfg.MODEL.ROI_HEADS.NUM_CLASSES = num_classes  # Update with the number of classes in your dataset

    # Solver configuration
    cfg.SOLVER.IMS_PER_BATCH = 2  # Adjust based on GPU memory
    cfg.SOLVER.BASE_LR = 0.00025
    cfg.SOLVER.MAX_ITER = 1500  # Adjust based on dataset size
    cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = 128

    # Output directory
    cfg.OUTPUT_DIR = output_dir
    os.makedirs(cfg.OUTPUT_DIR, exist_ok=True)

    return cfg

# Step 3: Train the model
def train_model(cfg):
    """
    Trains the model using the specified configuration.
    """
    trainer = DefaultTrainer(cfg)
    trainer.resume_or_load(resume=False)
    trainer.train()
    print(f"Training complete. Model saved to {cfg.OUTPUT_DIR}")

# Step 4: Visualize predictions
def test_and_visualize(cfg, image_path, output_path):
    """
    Tests the trained model on a new image and visualizes predictions.
    """
    from detectron2.engine import DefaultPredictor

    cfg.MODEL.WEIGHTS = os.path.join(cfg.OUTPUT_DIR, "model_final.pth")
    cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.3  # Adjust threshold
    predictor = DefaultPredictor(cfg)

    image = cv2.imread(image_path)
    outputs = predictor(image)

    v = Visualizer(image[:, :, ::-1], metadata=None, scale=1.2)
    out = v.draw_instance_predictions(outputs["instances"].to("cpu"))

    cv2.imwrite(output_path, out.get_image()[:, :, ::-1])
    print(f"Prediction saved to {output_path}")

def callback():
    print("Here we go")

# Example usage
if __name__ == "__main__":
    print("Training model")
    # Paths to dataset
    train_json = "dataset_coco.json"
    train_images = "data"

    # Output directory
    output_dir = "./output1"

    # Number of classes
    categories = [
    "door1",
    "window2",
    "window2",
    "window1",
    "door1",
    "door1",
    "table1",
    "armchair",
    "table1",
    "tub",
    "sink4",
    "sink3",
    "table1",
    "window2",
    "window2",
    "table2",
    "bed",
    "table1",
    "door1",
    "sofa2",
    "table1",
    "armchair",
    "sink3",
    "armchair",
    "armchair",
    "armchair",
    "table3"]  # Add all your classes
    num_classes = len(categories)

    # Register dataset
    register_dataset(train_json, train_images)

    # Configure and train the model
    cfg = get_training_config(output_dir, num_classes)
    train_model(cfg)

    # Test and visualize predictions
    # test_image = "1.jpg"
    # output_image = "output_prediction.jpg"
    # test_and_visualize(cfg, test_image, output_image)
    print("Model Trained. Here we go!")
    test_and_visualize(cfg,"./1.jpg", "out.jpg")
    # Gradio interface
    interface = gr.Interface(
        fn=callback,
        inputs=gr.Image(type="pil"),
        outputs=gr.Image(type="numpy"),
        title="Mask R-CNN Instance Segmentation",
        description="Upload an image for instance segmentation using Mask R-CNN."
    )

    interface.launch()