import cv2
import numpy as np
import sys
from pathlib import Path
import torch
from glob import glob

sys.path.append('..')
from torchlib.segneuralnet import SegmentationNeuralNet
from torchlib.post_processing_func import MAP_post
softmax = torch.nn.Softmax(1)
MAP = MAP_post()

def get_full_path(selected_model):
    if selected_model == 'Cross Entropy':
        model_path = r'models/selected/Segments_Seg1009_1.3.3_unetpad_cwce__adam_map_none_1_60_0/models/model_best.pth.tar'
    elif selected_model == 'Dice':
        model_path = r'models/selected/Segments_Seg1009_1.3.3_unetpad_cwdice__adam_map_none_1_60_0/models/model_best.pth.tar'
    elif selected_model == "J-REG":
        model_path = r'models/selected/Segments_Seg1009_1.3.3_unetpad_jreg__adam_map_none_1_60_0/models/model_best.pth.tar'
    else:
        assert False, f"MODEL {selected_model} NOT FOUD, OPT: 'Cross Entropy' 'Dice' 'J-REG'"
    return model_path


def load_model(full_url, use_cuda=True, gpu_id=0, verbose=False):
    full_url      = Path(full_url)
    full_path     = str(full_url)
    
    nameproject   = full_url.parent.parent.name
    patchproject  = full_url.parent.parent.parent
    ckpt_path     = '/'.join((full_url.parts)[-2:])
    file_name = full_url.name
    exp_type  = full_url.parent.parent.parent.name
    
    net = SegmentationNeuralNet(
        patchproject=patchproject, 
        nameproject=nameproject, 
        no_cuda=not use_cuda, parallel=False, seed=2021, 
        print_freq=False, gpu=gpu_id
    )

    check = net.load( full_path, verbose)
    
    if use_cuda:
        net.net.cuda(gpu_id)
    net.net.eval()
    
    return net


def get_results(srcs, net):

    def from2to3(src):
        zero = np.zeros(src.shape[:2])[..., None]
        return np.concatenate((zero, src), 2)

    inputs = torch.tensor(srcs)[None].float().cuda()
    outputs = net(inputs)
    output_soft = softmax(outputs).cpu().numpy()[0].transpose(1,2,0)
    outputs = outputs.cpu().numpy()[0].transpose(1,2,0)
    outputs = from2to3(outputs)
    output_soft = from2to3(output_soft)
    post = MAP(outputs, 2)[1]
    
    return output_soft, post

def resize_batchh(srcs):
    
    
        return srcs

def load_images(images_path, st_verbose=False):    
    srcs = np.array([cv2.imread(url, 0) for url in sorted(glob(f"{images_path}/*"))])
    len_srcs = len(srcs)
    if st_verbose:
        import streamlit as st
        st.text(f"Number of segmentations: {len_srcs}")

    if len_srcs != 60:
        np.random.seed(2021)
        ids = np.random.choice(list(range(len_srcs)), 60)
        
        if st_verbose:
            st.text("Warring!: Number of segmentations != 60, selecting 60 randomly")
            st.text(ids)
        srcs = srcs[ids]

    if srcs.max() > 1:
        srcs = srcs/srcs.max()
    
    return srcs
import cv2
import numpy as np
import sys
from pathlib import Path
import torch
from glob import glob

sys.path.append('..')
from torchlib.segneuralnet import SegmentationNeuralNet
from torchlib.post_processing_func import MAP_post
softmax = torch.nn.Softmax(1)
MAP = MAP_post()

def get_full_path(selected_model):
    if selected_model == 'Cross Entropy':
        model_path = r'models/selected/Segments_Seg1009_1.3.3_unetpad_cwce__adam_map_none_1_60_0/models/model_best.pth.tar'
    elif selected_model == 'Dice':
        model_path = r'models/selected/Segments_Seg1009_1.3.3_unetpad_cwdice__adam_map_none_1_60_0/models/model_best.pth.tar'
    elif selected_model == "J-REG":
        model_path = r'models/selected/Segments_Seg1009_1.3.3_unetpad_jreg__adam_map_none_1_60_0/models/model_best.pth.tar'
    else:
        assert False, f"MODEL {selected_model} NOT FOUD, OPT: 'Cross Entropy' 'Dice' 'J-REG'"
    return model_path


def load_model(full_url, use_cuda=True, gpu_id=0, verbose=False):
    full_url      = Path(full_url)
    full_path     = str(full_url)
    
    nameproject   = full_url.parent.parent.name
    patchproject  = full_url.parent.parent.parent
    ckpt_path     = '/'.join((full_url.parts)[-2:])
    file_name = full_url.name
    exp_type  = full_url.parent.parent.parent.name
    
    net = SegmentationNeuralNet(
        patchproject=patchproject, 
        nameproject=nameproject, 
        no_cuda=not use_cuda, parallel=False, seed=2021, 
        print_freq=False, gpu=gpu_id
    )

    check = net.load( full_path, verbose)
    
    if use_cuda:
        net.net.cuda(gpu_id)
    net.net.eval()
    
    return net


def get_results(srcs, net):

    def from2to3(src):
        zero = np.zeros(src.shape[:2])[..., None]
        return np.concatenate((zero, src), 2)

    inputs = torch.tensor(srcs)[None].float().cuda()
    outputs = net(inputs)
    output_soft = softmax(outputs).cpu().numpy()[0].transpose(1,2,0)
    outputs = outputs.cpu().numpy()[0].transpose(1,2,0)
    outputs = from2to3(outputs)
    output_soft = from2to3(output_soft)
    post = MAP(outputs, 2)[1]
    
    return output_soft, post

def resize_batchh(srcs):
    
    
        return srcs

def load_images(images_path, st_verbose=False):    
    srcs = np.array([cv2.imread(url, 0) for url in sorted(glob(f"{images_path}/*"))])
    len_srcs = len(srcs)
    if st_verbose:
        import streamlit as st
        st.text(f"Number of segmentations: {len_srcs}")

    if len_srcs != 60:
        np.random.seed(2021)
        ids = np.random.choice(list(range(len_srcs)), 60)
        
        if st_verbose:
            st.text("Warring!: Number of segmentations != 60, selecting 60 randomly")
            st.text(ids)
        srcs = srcs[ids]

    if srcs.max() > 1:
        srcs = srcs/srcs.max()
    
    return srcs