import os
import numpy as np
import tensorflow.compat.v1 as tf
tf.disable_v2_behavior()
from PIL import Image
import imageio
from net import Network
from utils.rgb_ind_convertor import ind2rgb, floorplan_fuse_map

class DeepFloorPlanModel:
    def __init__(self, model_dir='pretrained', input_size=(512, 512)):
        self.input_size = input_size
        self.model_dir = model_dir
        self._build_graph()
        self._load_weights()

    def _build_graph(self):
        tf.compat.v1.reset_default_graph()
        self.sess = tf.compat.v1.Session()
        self.x = tf.compat.v1.placeholder(shape=[1, self.input_size[0], self.input_size[1], 3], dtype=tf.float32, name='inputs')
        self.network = Network()
        logits1, logits2 = self.network.forward(self.x, init_with_pretrain_vgg=False)
        self.rooms = self.network.convert_one_hot_to_image(logits1, act='softmax', dtype='int')
        self.close_walls = self.network.convert_one_hot_to_image(logits2, act='softmax', dtype='int')
        self.sess.run(tf.compat.v1.global_variables_initializer())
        self.sess.run(tf.compat.v1.local_variables_initializer())
        self.saver = tf.compat.v1.train.Saver()

    def _load_weights(self):
        ckpt = tf.train.latest_checkpoint(self.model_dir)
        if ckpt is None:
            print(f"[ERROR] No checkpoint found in {self.model_dir}")
            raise FileNotFoundError(f"No checkpoint found in {self.model_dir}")
        print(f"[INFO] Restoring model weights from {ckpt}")
        self.saver.restore(self.sess, ckpt)

    def predict(self, image):
        # Accepts a numpy array or PIL image, returns a numpy array (segmentation mask)
        if isinstance(image, Image.Image):
            image = np.array(image)
        if image.shape[-1] == 4:
            image = image[..., :3]
        im_resized = np.array(Image.fromarray(image).resize(self.input_size, Image.BICUBIC)) / 255.0
        im_resized = im_resized.astype(np.float32)
        im_resized = np.reshape(im_resized, (1, self.input_size[0], self.input_size[1], 3))
        out1, out2 = self.sess.run([self.rooms, self.close_walls], feed_dict={self.x: im_resized})
        out1 = np.squeeze(out1)
        out2 = np.squeeze(out2)
        # Merge logic: set out1 pixels to 9/10 where out2==1/2
        out1[out2==2] = 10
        out1[out2==1] = 9
        # Convert to RGB for visualization
        out_rgb = ind2rgb(out1, color_map=floorplan_fuse_map)
        out_rgb = out_rgb.astype(np.uint8)
        return out_rgb

    def close(self):
        self.sess.close()