import web

urls=(
	'/','gradeList',
	'/add','gradeAdd',
	'/del/(.+)','gradeDel',
	'/edit/(.+)','gradeEdit',
	'/editSave','gradeEditSave',
	)

app=web.application(urls,globals())

render=web.template.render('templates/')
db=web.database(dbn='sqlite',db='datafile/grade.db')

class gradeList:
	def GET(self):
		grades=db.select('grade')
		return render.list(grades)
class gradeAdd:
	def GET(self):
		return render.add()

	def POST(self):
		i=web.input()
		grade=db.insert('grade',name=i.gradeName,value=i.gradeValue,id=i.gradeId)
		raise web.seeother('/')

class gradeDel:
	def GET(self,id):
		db.delete('grade',where ='id=$id',vars=locals())
		raise web.seeother('/')

class gradeEdit:
	def GET(self,id):
		grade=db.select('grade',where='id=$id',vars=locals())[0]

		sa=[]
		for gg in grade:
			sa.append(gg)
		print('edit~grade',sa)
		return render.edit(grade)

class gradeEditSave:
	def POST(self):
		i=web.input()
		grade=db.update('grade',name=i.gradeName,value=i.gradeValue,where ='id=$gradeId',vars=i)
		raise web.seeother('/')

if __name__ == '__main__':
	app.run()


from __future__ import division
from builtins import range
import numpy as np
from functools import wraps
from nose.plugins.attrib import attr
import os
import matplotlib.pyplot as plt

from pyhsmm import models, distributions
from pybasicbayes.util import testing

##########
#  util  #
##########

def runmultiple(n):
    def dec(fn):
        @wraps(fn)
        def wrapper():
            fig = plt.figure()
            for i in range(n):
                yield fn, fig
            plt.close('all')
        return wrapper
    return dec

def mkdir(path):
    # from
    # http://stackoverflow.com/questions/600268/mkdir-p-functionality-in-python
    import errno
    try:
        os.makedirs(path)
    except OSError as exc:
        if exc.errno == errno.EEXIST and os.path.isdir(path):
            pass
        else: raise

figure_dir_path = os.path.join(os.path.dirname(__file__),'figures')
mkdir(figure_dir_path)

###########
#  tests  #
###########

@attr('slow')
@runmultiple(2)
def discrete_geweke_test(fig):
    Nstates = 2
    Nemissions = 2
    alpha = 3.
    init_state_concentration=3.
    T = 10
    num_iter = 10000
    num_checks = 10

    obs_distns = [distributions.Categorical(K=Nemissions,alpha_0=1.)
            for _ in range(Nstates)]

    hmm = models.HMM(
            alpha=alpha,init_state_concentration=init_state_concentration,
            obs_distns=obs_distns)

    # generate state sequences and parameters from the prior
    prior_stateseqs = []
    prior_weights = []
    for itr in range(num_iter):
        hmm.resample_model() # sample parameters from the prior
        _, stateseq = hmm.generate(T,keep=False)
        prior_stateseqs.append(stateseq)
        prior_weights.append(hmm.obs_distns[0].weights)
    prior_stateseqs = np.array(prior_stateseqs)
    prior_weights = np.array(prior_weights)

    # generate state sequences and parameters using Gibbs
    hmm.generate(T,keep=True)
    s = hmm.states_list[0]

    gibbs_stateseqs = []
    gibbs_weights = []
    for itr in range(num_iter):
        s.data = None
        hmm._generate_obs(s)  # resamples data given state sequence, obs params
        hmm.resample_model()  # resamples everything else as usual
        gibbs_stateseqs.append(s.stateseq)
        gibbs_weights.append(hmm.obs_distns[0].weights)
    gibbs_stateseqs = np.array(gibbs_stateseqs)
    gibbs_weights = np.array(gibbs_weights)

    # test that they look similar by checking probability of co-assignment
    time_indices = np.arange(T)
    for itr in range(num_checks):
        i,j = np.random.choice(time_indices,replace=False,size=2)
        prior_prob_of_coassignment = (prior_stateseqs[:,i] == prior_stateseqs[:,j]).std()
        gibbs_prob_of_coassignment = (gibbs_stateseqs[:,i] == gibbs_stateseqs[:,j]).std()

        assert np.isclose(
                prior_prob_of_coassignment,gibbs_prob_of_coassignment,
                rtol=0.025,atol=0.025,
                )

    # test that they look similar by checking parameters
    testing.populations_eq_quantile_plot(prior_weights,gibbs_weights,fig=fig)
    figpath = os.path.join(figure_dir_path,'discrete_geweke_test_weights.pdf')
    plt.savefig(figpath)

@attr('slow')
@runmultiple(2)
def discrete_geweke_multiple_seqs_test(fig):
    Nstates = 2
    Nemissions = 2
    alpha = 3.
    init_state_concentration=3.
    T = 10
    num_seqs = 3
    num_iter = 10000
    num_checks = 10

    obs_distns = [distributions.Categorical(K=Nemissions,alpha_0=1.)
            for _ in range(Nstates)]

    hmm = models.HMM(
            alpha=alpha,init_state_concentration=init_state_concentration,
            obs_distns=obs_distns)

    # generate state sequences and parameters from the prior
    prior_stateseqss = [[] for _ in range(num_seqs)]
    prior_weights = []
    for itr in range(num_iter):
        hmm.resample_model() # sample parameters from the prior

        for itr2 in range(num_seqs):
            _, stateseq = hmm.generate(T,keep=False)
            prior_stateseqss[itr2].append(stateseq)

        prior_weights.append(hmm.obs_distns[0].weights)

    prior_stateseqss = np.array(prior_stateseqss)
    assert prior_stateseqss.shape == (num_seqs,num_iter,T)
    prior_weights = np.array(prior_weights)

    # generate state sequences and parameters using Gibbs
    for itr in range(num_seqs):
        hmm.generate(T,keep=True)
    assert len(hmm.states_list) == num_seqs

    gibbs_stateseqss = [[] for _ in range(num_seqs)]
    gibbs_weights = []
    for itr in range(num_iter):
        for s in hmm.states_list:
            s.data = None
            hmm._generate_obs(s)  # resamples data given state sequence, obs params
        hmm.resample_model()  # resamples everything else as usual

        for itr2, s in enumerate(hmm.states_list):
            gibbs_stateseqss[itr2].append(s.stateseq)

        gibbs_weights.append(hmm.obs_distns[0].weights)

    gibbs_stateseqss = np.array(gibbs_stateseqss)
    assert gibbs_stateseqss.shape == (num_seqs,num_iter,T)
    gibbs_weights = np.array(gibbs_weights)

    # test that they look similar by checking probability of co-assignment
    time_indices = np.arange(T)
    seq_indices = np.arange(num_seqs)
    for itr in range(num_checks):
        i,j = np.random.choice(time_indices,replace=False,size=2)
        si,sj = np.random.choice(seq_indices,replace=True,size=2)
        prior_prob_of_coassignment = \
                (prior_stateseqss[si,:,i] == prior_stateseqss[sj,:,j]).std()
        gibbs_prob_of_coassignment = \
                (gibbs_stateseqss[si,:,i] == gibbs_stateseqss[sj,:,j]).std()

        assert np.isclose(
                prior_prob_of_coassignment,gibbs_prob_of_coassignment,
                rtol=0.025,atol=0.025,
                )

    # test that they look similar by checking parameters
    testing.populations_eq_quantile_plot(prior_weights,gibbs_weights,fig=fig)
    figpath = os.path.join(figure_dir_path,
            'discrete_geweke_multiple_seqs_test_weights.pdf')
    plt.savefig(figpath)


#!/usr/bin/env python3

#### IMPORTS ####

import rospy
from nav_msgs.msg import Odometry
from tf.transformations import euler_from_quaternion
from geometry_msgs.msg import PoseWithCovarianceStamped, Pose, Twist, Vector3

#___________________________________________________________________________________________________

#### CALLBACK-FUNCTION ####

def callback_hector_pose(hec_pose):

        global last_timestamp, last_rotat, x_last, y_last, theta_last

        if hec_pose.header.stamp > last_timestamp:

                #calculate euler from quaternion
                euler_rotat = euler_from_quaternion([hec_pose.pose.pose.orientation.x,
                                                        hec_pose.pose.pose.orientation.y,
                                                        hec_pose.pose.pose.orientation.z,
                                                        hec_pose.pose.pose.orientation.w])[2]

                #calculate time difference
                time_diff = (hec_pose.header.stamp - last_timestamp).to_sec()
                #calculate velocitys
                x_vel = ((hec_pose.pose.pose.position.x - x_last)/time_diff)
                y_vel = ((hec_pose.pose.pose.position.y - y_last)/time_diff)
                theta_vel = ((euler_rotat - last_rotat)/time_diff)

                #pass hec_pose.pose.position to x_last and y_last for velocity calculation
                x_last = hec_pose.pose.pose.position.x
                y_last = hec_pose.pose.pose.position.y
                #pass rotat in euler to last_rotat for velocity calculation
                last_rotat = euler_rotat
                #pass hec_pose.header.stamp to last_timestamp
                last_timestamp = hec_pose.header.stamp

                #fill the odom_msg for publishing the odom from hector topic
                odom_msg.header.stamp = rospy.Time.now()  #timestamp for the published msg
                odom_msg.pose.pose = hec_pose.pose.pose

                #checking if the pose msg should use a static covariance msg or the one passed from hector mapping
                if use_static_pose_covariance:
                        odom_msg.pose.covariance = static_pose_covariance
                else:
                        odom_msg.pose.covariance = hec_pose.pose.covariance
                odom_msg.twist.twist = Twist(Vector3(x_vel, y_vel, 0), Vector3(0, 0, theta_vel))
                odom_msg.twist.covariance = static_twist_covariance

                #publish odom from hector topic
                odom_pub.publish(odom_msg)

                #define sleep_rate
                rate = rospy.Rate(sleep_rate)
                #sleep
                rate.sleep()

#___________________________________________________________________________________________________

#### MAIN-CODE ####

if __name__ == '__main__':
        #initiate ROS node
        rospy.init_node("hector_to_odom_topic_publisher")

        #get parameters from rosparam server
        node_name = rospy.get_name()
        parent_frame = rospy.get_param(node_name + "/parent_frame", "odom")
        child_frame = rospy.get_param(node_name + "/child_frame", "base_link")
        odom_topic_name = rospy.get_param(node_name + "/odom_topic_name", "odom/hector")
        sleep_rate = rospy.get_param(node_name + "/sleep_rate", 4)

        #decide if you want to use static pose covariance
        use_static_pose_covariance = rospy.get_param(node_name + "/use_static_pose_covariance", False)

        #only change the static covariance value for the main axis so
        #you don't have to give a whole covariance matrix with 36 values
        static_pose_covariance_val = rospy.get_param(node_name + "/static_pose_covariance_val", 0.2)
        static_twist_covariance_val = rospy.get_param(node_name + "/static_twist_covariance_val", 0.4)

        #if needed you can change the whole static_covariance matrix through rosparam
        #if not needed it defaults to the values seen below
        static_pose_covariance = rospy.get_param(node_name + "/static_pose_covariance",
                                                [static_pose_covariance_val, 0, 0, 0, 0, 0,
                                                0, static_pose_covariance_val, 0, 0, 0, 0,
                                                0, 0, 0, 0, 0, 0,
                                                0, 0, 0, 0, 0, 0,
                                                0, 0, 0, 0, 0, 0,
                                                0, 0, 0, 0, 0, static_pose_covariance_val])

        static_twist_covariance = rospy.get_param(node_name + "/static_twist_covariance",
                                                  [static_twist_covariance_val, 0, 0, 0, 0, 0,
                                                   0, static_twist_covariance_val, 0, 0, 0, 0,
                                                   0, 0, 0, 0, 0, 0,
                                                   0, 0, 0, 0, 0, 0,
                                                   0, 0, 0, 0, 0, 0,
                                                   0, 0, 0, 0, 0, static_twist_covariance_val])

        #define an odometry publisher
        odom_pub = rospy.Publisher(odom_topic_name, Odometry, queue_size=25)

        #instance of an odometry message
        odom_msg = Odometry()
        odom_msg.header.frame_id = parent_frame
        odom_msg.child_frame_id = child_frame

        #definde variables needed for publishing odom topic
        x_vel, y_vel, theta_vel = 0.0, 0.0, 0.0 #velocity
        last_timestamp = rospy.Time.now()       #save last time that we received a tf
        x_last, y_last = 0.0, 0.0               #save last translation
        last_rotat = 0.0                        #save last rotation in euler

        #subscribe to hector poseupdate
        rospy.Subscriber("poseupdate", PoseWithCovarianceStamped, callback_hector_pose, queue_size = 1)

        rospy.spin()

#___________________________________________________________________________________________________


import torch
import torchvision
from torchvision.models.detection import fasterrcnn_resnet50_fpn
from torchvision.datasets import VOCDetection
from torch.utils.data import DataLoader
import json
from utils import object_dict, transform, collate_fn
from caculate_mAP import calculate_mAP

# 定义数据加载器 train_loader 和 test_loader
my_device = 'cuda'

# 数据集路径
root_dir = '/Users/yanxinyu/pythonProject/HW2_FasterRCNN'
# 数据集使用的年份
year = '2012'
# 训练/测试集
train_set = VOCDetection(root=root_dir, year=year, image_set='train', download=False)
val_set = VOCDetection(root=root_dir, year=year, image_set='val', download=False)

# 将数据集包装为 DataLoader，batch_size可以根据计算资源调整
train_loader = DataLoader(train_set, batch_size=16, shuffle=True, num_workers=4, collate_fn=collate_fn)
test_loader = DataLoader(val_set, batch_size=16, shuffle=False, num_workers=4, collate_fn=collate_fn)

# 创建 Faster R-CNN 模型
model = fasterrcnn_resnet50_fpn(num_classes=21, weights_backbone=torchvision.models.ResNet50_Weights.IMAGENET1K_V1)

# 定义优化器
optimizer = torch.optim.SGD(model.parameters(), lr=0.01, momentum=0.1)
lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=1, gamma=0.8)

def estimate_loss(model, data_loader, sample_ratio=0.05):
    sample_size = int(len(data_loader.dataset) * sample_ratio)
    sampled_data_loader = torch.utils.data.DataLoader(data_loader.dataset, batch_size=data_loader.batch_size, sampler=torch.utils.data.RandomSampler(data_loader.dataset, replacement=True, num_samples=sample_size), collate_fn=collate_fn)
    total_loss = 0
    with torch.no_grad():
        for images, targets in sampled_data_loader:
            images = list(image.to(my_device) for image in images)
            targets = [{k: v.to(my_device) if torch.is_tensor(v) else v for k, v in t.items()} for t in targets]
            model.to(my_device)
            loss_dict = model(images, targets)
            total_loss += sum(loss for loss in loss_dict.values()).item()
    estimated_loss = total_loss / sample_size * len(data_loader.dataset)
    return estimated_loss

def estimate_mAP(model, data_loader, iou_threshold=0.5, conf_threshold=0.05, device='cpu', sample_ratio=0.05):
    sample_size = int(len(data_loader.dataset) * sample_ratio)
    sampled_data_loader = torch.utils.data.DataLoader(data_loader.dataset, batch_size=data_loader.batch_size, sampler=torch.utils.data.RandomSampler(data_loader.dataset, replacement=True, num_samples=sample_size), collate_fn=collate_fn)
    return calculate_mAP(model, sampled_data_loader, iou_threshold, conf_threshold, device)

if __name__ == '__main__':
# 迭代训练
    num_epochs = 20
    loss_and_mAP = {
            'train_loss':[],
            'test_loss':[],
            'mAP':[]
        }

    # 计算初始值的loss
    train_loss = estimate_loss(model, train_loader, sample_ratio=0.2)
    test_loss = estimate_loss(model, test_loader, sample_ratio=0.2)
    print(f"epoch={0}, round={0}, train_loss={train_loss}")
    print(f"epoch={0}, round={0}, test_loss={test_loss}")
    print('lr=',optimizer.param_groups[0]['lr'])
    loss_and_mAP['train_loss'].append(train_loss)
    loss_and_mAP['test_loss'].append(test_loss)

    for epoch in range(num_epochs):
        num = 0
        for images, targets in train_loader:
            # 将图像和目标数据加载到设备
            images = [image.to(my_device) for image in images]
            targets = [{k: v.to(my_device) if torch.is_tensor(v) else v for k, v in t.items()} for t in targets]
            model.to(my_device)

            # 设置模型为训练模式
            model.train()

            # 清零优化器的梯度
            optimizer.zero_grad()

            # 获取模型的预测结果和损失
            loss_dict = model(images, targets)
            print(loss_dict)

            # 计算总损失
            loss = sum(loss for loss in loss_dict.values())

            # 反向传播和梯度更新
            loss.backward()
            optimizer.step()
            # 输出当前迭代的损失函数值
            num += 1
            print(f"Epoch [{epoch+1}/{num_epochs}], Loss: {loss_dict}")
            print(f"Round [{num}], Loss: {loss}")

            if num%10==0:
                train_loss = estimate_loss(model, train_loader, sample_ratio=0.2)
                test_loss = estimate_loss(model, test_loader, sample_ratio=0.2)
                print(f"epoch={epoch}, round={num}, train_loss={train_loss}")
                print(f"epoch={epoch}, round={num}, test_loss={test_loss}")
                print('lr=',optimizer.param_groups[0]['lr'])
                loss_and_mAP['train_loss'].append(train_loss)
                loss_and_mAP['test_loss'].append(test_loss)

                # 打印结果
                print(f'Epoch {epoch + 1}/{num_epochs}, round {num}'
                      f'Train Loss: {train_loss:.3f}, '
                      f'Test Loss: {test_loss:.3f}, '
                )
                with open('loss_and_mAP.json', 'w') as file:
                    json.dump(loss_and_mAP, file)

        lr_scheduler.step()
        mAP = estimate_mAP(model, test_loader, sample_ratio=0.10)
        loss_and_mAP['mAP'].append(mAP)
        torch.save(model.state_dict(), f'new_model_epoch{epoch+1}.pth')


import pytest
import importlib
import opentrons
from typing import Any, TYPE_CHECKING

if TYPE_CHECKING:
    from opentrons_shared_data.deck.dev_types import RobotModel

@pytest.fixture(autouse=True)
def reload_module(robot_model: "RobotModel") -> None:
    importlib.reload(opentrons.calibration_storage)

@pytest.fixture
def starting_calibration_data(
    robot_model: "RobotModel", ot_config_tempdir: Any
) -> None:
    """
    Starting calibration data fixture.

    Adds dummy data to a temporary directory to test delete commands against.
    """
    from opentrons.calibration_storage import save_robot_deck_attitude

    if robot_model == "OT-3 Standard":
        save_robot_deck_attitude([[1, 0, 0], [0, 1, 0], [0, 0, 1]], "pip1")
    else:
        save_robot_deck_attitude([[1, 0, 0], [0, 1, 0], [0, 0, 1]], "pip1", "mytiprack")

def test_save_deck_attitude(ot_config_tempdir: Any, robot_model: "RobotModel") -> None:
    """
    Test saving deck attitude calibrations.
    """
    from opentrons.calibration_storage import (
        get_robot_deck_attitude,
        save_robot_deck_attitude,
    )

    assert get_robot_deck_attitude() is None
    if robot_model == "OT-3 Standard":
        save_robot_deck_attitude([[1, 0, 0], [0, 1, 0], [0, 0, 1]], "pip1")
    else:
        save_robot_deck_attitude([[1, 0, 0], [0, 1, 0], [0, 0, 1]], "pip1", "mytiprack")
    assert get_robot_deck_attitude() != {}

def test_get_deck_calibration(
    starting_calibration_data: Any, robot_model: "RobotModel"
) -> None:
    """
    Test ability to get a deck calibration model.
    """
    from opentrons.calibration_storage import get_robot_deck_attitude

    # needed for proper type checking unfortunately
    from opentrons.calibration_storage.ot3.models.v1 import (
        DeckCalibrationModel as OT3DeckCalModel,
    )
    from opentrons.calibration_storage.ot2.models.v1 import (
        DeckCalibrationModel as OT2DeckCalModel,
    )

    robot_deck = get_robot_deck_attitude()
    if robot_model == "OT-3 Standard":
        assert robot_deck == OT3DeckCalModel(
            attitude=[[1, 0, 0], [0, 1, 0], [0, 0, 1]],
            lastModified=robot_deck.lastModified,
            source=robot_deck.source,
            pipetteCalibratedWith="pip1",
        )
    else:
        assert robot_deck == OT2DeckCalModel(
            attitude=[[1, 0, 0], [0, 1, 0], [0, 0, 1]],
            last_modified=robot_deck.last_modified,
            source=robot_deck.source,
            pipette_calibrated_with="pip1",
            tiprack="mytiprack",
        )

def test_delete_deck_calibration(starting_calibration_data: Any) -> None:
    """
    Test delete deck calibration.
    """
    from opentrons.calibration_storage import (
        get_robot_deck_attitude,
        delete_robot_deck_attitude,
    )

    assert get_robot_deck_attitude() != {}
    assert get_robot_deck_attitude().attitude == [[1, 0, 0], [0, 1, 0], [0, 0, 1]]
    delete_robot_deck_attitude()
    assert get_robot_deck_attitude() is None


import pytest

from wattson.data import instance_data, region_data
from wattson.estimator import (
    calculate_scope_2,
    calculate_scope_3,
    compute_load_to_kwh,
    estimate_carbon_emissions,
)

SAMPLE_INSTANCE = instance_data["db.m4.10xlarge"]
SAMPLE_REGION = region_data["eu-north-1"]

def test_calculate_scope_2() -> None:
    with_PUE = 245.6 * 1.2
    expected_result = with_PUE / 1000 * 1 * 8
    result = calculate_scope_2(
        instance=SAMPLE_INSTANCE, region=SAMPLE_REGION, hours=1, load=50
    )
    assert result == expected_result
    # Cross-validate with a known example from Teads calculator
    assert round(result, 1) == 2.4

def test_calculate_scope_3() -> None:
    result = calculate_scope_3(instance=SAMPLE_INSTANCE, hours=1)
    # Cross-validate with a known example from Teads calculator
    assert result == 34.5

def test_lerp() -> None:
    assert round(compute_load_to_kwh(SAMPLE_INSTANCE.load_to_kwh, 25)) == 190

@pytest.mark.parametrize(
    ("region", "has_compensation", "expected_emissions"),
    [
        ("ca-central-1", True, 14.1),
        ("eu-north-1", False, 11.3),
        ("ap-northeast-1", False, 22.8),
    ],
)
def test_complete_estimate(
    region: str, has_compensation: bool, expected_emissions: float
) -> None:
    instance_type = "a1.large"
    hours = 3
    load = 50

    result = estimate_carbon_emissions(
        instance_type=instance_type,
        hours=hours,
        load_percentage=load,
        region=region,
    )

    # Cross-validate with a known example from Teads calculator
    assert round(result.scope_2_co2eq + result.scope_3_co2eq, 1) == expected_emissions
    assert (
        "has been compensated by renewable energy credits" in result.details
    ) == has_compensation
    assert result.compensated == has_compensation

# fmt: off
@pytest.mark.parametrize("value_set", [
    dict(region="xx-mars", instance_type="db.m4.10xlarge", hours=1),
    dict(region="eu-north-1", instance_type="deep.thought", hours=1),
    dict(region="eu-north-1", instance_type="db.m4.10xlarge", hours=-1),
    dict(region="eu-north-1", instance_type="db.m4.10xlarge", hours=1, load_percentage=1500),  # noqa: E501
])
# fmt: on
def test_nonsensical_values(value_set: dict) -> None:  # type: ignore[type-arg]
    with pytest.raises(ValueError):
        estimate_carbon_emissions(**value_set)


import io
import itertools
import logging
import os
import re
import shutil
import glob
import fnmatch
from typing import List, Tuple
from tqdm import tqdm
import cv2
import google.auth
import numpy as np
import pandas as pd
import torch
import torch.nn as nn
from google.cloud import vision
from konlpy.tag import Komoran

from analyzer.constants import IMAGE_SIZE
from analyzer.data_utils import read_image
from analyzer.utils import to_best_device
from analyzer.v2.datasets import EvalSegmentationSubsDataset
from analyzer.v2.model import AttentionSubsBoxerModel

logger = logging.getLogger(__name__)
GOOGLE_MAX_HEIGHT = 2050
GOOGLE_MAX_WIDTH = 1536
NB_ROWS = 10
NB_COLUMNS = 3
COLUMN_HEIGHT = int(GOOGLE_MAX_HEIGHT / NB_ROWS)
COLUMN_WIDTH = int(GOOGLE_MAX_WIDTH / NB_COLUMNS)

BREAKS = vision.TextAnnotation.DetectedBreak.BreakType

class obj:
    # constructor
    def __init__(self, dict1):
        self.__dict__.update(dict1)

class VideoAnalyzer:

    def __init__(self, work_directory: str, model_name: str, target_directory: str,storage, skip_frames: int = 30,
                 bb_margin: int = 10):
        self.work_directory = work_directory
        self.target_directory = target_directory
        self.skip_frames = int(skip_frames)
        self.ensure_dir(work_directory)
        self.ensure_dir(target_directory)
        self.threshold = 0.5
        self.bb_margin = bb_margin
        self.vision_client = vision.ImageAnnotatorClient()
        self.storage = storage
        self.model = self._get_model(model_name)
        self.treated = set([])
        logger.info("Loading Komoran...")
        self.analyzer = Komoran()
        logger.info("...Komoran loaded")
        logger.info(f"Working with directory {work_directory}")
        logger.info(f"Target       directory {target_directory}")
        logger.info(f"Loaded model           {model_name}")

    def ensure_dir(self, file_path):
        os.makedirs(file_path, exist_ok=True)

    def treat_incoming_file(self, file_path: str) -> List[str]:
        logger.info(f"Treating file {file_path}")
        work_file_path = self.prepare_file(file_path)
        if work_file_path:
            prefix_splitted: str = self.split_file(work_file_path)
            annotation_file, byproducts = self.create_annotations(prefix_splitted)
            annotation_file_with_subs = self.add_subs(annotation_file, prefix_splitted)
            annotation_file_polished = self.polish(annotation_file_with_subs, prefix_splitted, work_file_path)
            annotation_file_final = self.parse_korean(annotation_file_polished, prefix_splitted)
            products = [annotation_file_final, work_file_path, file_path]
            products = self.move_products(prefix_splitted, products)
            # Get useful by products
            if len(byproducts) > 0:
                products.extend(byproducts)
            logger.info(f"Done ! {file_path} treated")
            return products
        return []

    def move_products(self, prefix: str, products_path: [str]) -> List[str]:
        logger.info(f"Moving products of {prefix} to {self.target_directory}")
        dest_dir = f"{self.target_directory}/{prefix}"
        self.ensure_dir(dest_dir)
        final_paths: List[str] = []
        for path in products_path:
            file_name = os.path.basename(path)
            final_path = f"{dest_dir}/{file_name}"
            shutil.move(path, final_path)
            logger.info(f"------ {path} moved to {final_path}")
            final_paths.append(final_path)
        return final_paths

    def parse_korean(self, annotation_file_path: str, prefix: str) -> str:
        logger.info("Parsing korean subs")
        final_annotation_file_path = f"{self.work_directory}/{prefix}.csv"
        df_in = pd.read_csv(annotation_file_path)
        extends = []
        nb_lines = len(df_in)
        for i, row in df_in.iterrows():
            parsed = self.analyzer.pos(row['subs'])
            extension = []
            for p in parsed:
                extension.append(p[0])
                extension.append(p[1])
            extends.append(extension)
            if (i + 1) % 100 == 0:
                logger.debug(f"{i + 1} / {nb_lines} subs analyzed")

        extends = np.array(list(zip(*itertools.zip_longest(*extends, fillvalue=''))))
        nb_extra_columns = len(extends[0])
        for i in range(nb_extra_columns):
            df_in[f"parsed_{i}"] = extends[:, i]
        df_in.to_csv(final_annotation_file_path, encoding='utf-8', index=False)
        return final_annotation_file_path

    def polish(self, annotation_file_path: str, prefix: str, video_file_path: str):
        logger.info("Polishing of the file")
        final_annotation_file_path = f"{self.work_directory}/{prefix}-polished.csv"
        if os.path.exists(final_annotation_file_path):
            logger.info("--- Polishing already done")
        else:
            df_in = pd.read_csv(annotation_file_path)
            video = cv2.VideoCapture(video_file_path)
            fps = video.get(cv2.CAP_PROP_FPS)
            spf = 1 / fps
            spf_for_sampling_rate = self.skip_frames * spf

            # Merge lines with same subs
            data = []
            curr_subs = None
            curr_frame_start, curr_frame_end = None, None
            for i, annotation in df_in.iterrows():
                start = annotation['start']
                end = annotation['end']
                subs = annotation['subs']
                if curr_subs == subs:
                    curr_frame_end = end
                else:
                    if not curr_subs is None:
                        data.append([curr_subs, curr_frame_start, curr_frame_end])
                    if pd.isna(subs):
                        curr_frame_start = None
                        curr_frame_end = None
                        curr_subs = None
                    else:
                        curr_frame_start = start
                        curr_frame_end = end
                        curr_subs = subs
            if not curr_subs is None:
                data.append([curr_subs, curr_frame_start, curr_frame_end])
            df = pd.DataFrame(columns=['subs', 'start', 'end'], data=data)
            df['start'] = df['start'] * spf_for_sampling_rate
            df['end'] = df['end'] * spf_for_sampling_rate
            df.to_csv(final_annotation_file_path, index=False)
        return final_annotation_file_path

    def prepare_file(self, file_path: str) -> str:
        filename_without_extension = os.path.splitext(os.path.basename(file_path))[0]
        out_file = f"{self.work_directory}/{filename_without_extension}.mp4"
        if os.path.exists(out_file):
            print("File already converted")
            return out_file
        else:
            print(f"Converting {file_path} into {out_file}")
            os.system(f"ffmpeg -i {file_path} {out_file}")
            if os.path.isfile(out_file):
                print(f"Conversion successful")
                return out_file
            else:
                print(f"Could not convert file")
                return None

    def split_file(self, file_path) -> str:
        """
        Extract frames of the video at constant rate and output them to the work directory
        :param file_path: Path of the video
        :return: Prefix of the video file
        """
        print(f"Splitting file {file_path}")
        filename_without_extension = os.path.splitext(os.path.basename(file_path))[0]
        if os.path.exists(f"{self.work_directory}/{filename_without_extension}-0.jpg"):
            print(f"--- File already splitted")
        else:
            cap = cv2.VideoCapture(file_path)
            i = 0
            stop_all = False
            while cap.isOpened() and not stop_all:
                for j in range(0, self.skip_frames):
                    ret, frame = cap.read()
                    if ret == False:
                        stop_all = True
                if not stop_all:
                    cv2.imwrite(f"{self.work_directory}/{filename_without_extension}-{str(i)}.jpg", frame)
                    i += 1
                    if i % 10 == 0:
                        logger.info(f"--- Exported {i} frames")
            cap.release()
            cv2.destroyAllWindows()
            logger.info(f"End splitting file {file_path} into {self.work_directory} with prefix {filename_without_extension}")
        return filename_without_extension

    def create_annotations(self, prefix_splitted: str) -> Tuple[str, List[str]]:
        """
        Analyze splitted files to detect bounding boxes
        :param prefix_splitted: Prefix of the files containing the frames of the video
        :return: Path of the file containing the following columns
        'filename', 'start', 'end'
        filename: string, name of the file with the subs
        start: int, index of the continuous series of frames containing the same subs
        end: int, index of the continuous series of frames containing the same subs
        """
        logger.info("Predicting subs from video")
        annotations_file_path = f"{self.work_directory}/annotations_{prefix_splitted}.csv"
        byproducts: List[str] = []
        if os.path.exists(annotations_file_path):
            logger.info("--- subs prediction already done")
        else:
            size = (IMAGE_SIZE, IMAGE_SIZE)
            data = []
            should_analyze_file = True
            start = -1
            i = 0
            x0 = -1
            x1 = -1
            y0 = -1
            y1 = -1
            current_file_path: str = None
            current_zone: np.array = None
            num_files = len(fnmatch.filter(os.listdir(self.work_directory), f"{prefix_splitted}-*.jpg"))
            for i in tqdm(range(num_files)):
                splitted_file = f"{self.work_directory}/{prefix_splitted}-{i}.jpg"
                if os.path.exists(splitted_file):
                    im = read_image(splitted_file)
                    im = cv2.resize(im, size)
                    if not should_analyze_file:
                        # Check if the bounding boxes contain the same information
                        image_zone = np.array(im[y0: y1, x0: x1])
                        if not self.same_subs(current_zone, image_zone):
                            # This is the end of the optimization for this sub
                            should_analyze_file = True
                            data.extend([[current_file_path, start, i - 1]])
                            current_zone = None
                            current_file_path = None
                            start = None
                        else:
                            logger.debug(f"Optimize frame {i} with frame {start}")
                    if should_analyze_file:
                        resized_file_path = f"{self.work_directory}/{prefix_splitted}-resized-{i}.jpg"
                        cv2.imwrite(resized_file_path, cv2.cvtColor(im, cv2.COLOR_RGB2BGR))
                        test_ds = EvalSegmentationSubsDataset(pd.DataFrame([{'filename': resized_file_path}]))
                        x = test_ds[0]
                        xx = to_best_device(torch.FloatTensor(x[None,]))
                        predicted_bbs = self.model(xx)
                        x0, y0, x1, y1, present = predicted_bbs.detach().cpu()[0]
                        has_bb: bool = present >= self.threshold
                        if has_bb:
                            y0 = int(np.floor(max(y0 - self.bb_margin, 0)))
                            y1 = int(np.ceil(min(y1 + self.bb_margin, IMAGE_SIZE - 1)))
                            x0 = int(np.floor(max(x0 - self.bb_margin, 0)))
                            x1 = int(np.ceil(min(x1 + self.bb_margin, IMAGE_SIZE - 1)))
                            has_bb = self._is_trustable_bounding_box(x0, y0, x1, y1)
                            if not has_bb:
                                postmortem_file = f"{self.work_directory}/{prefix_splitted}-postmortem-{i}.jpg"
                                cv2.imwrite(postmortem_file, cv2.cvtColor(im, cv2.COLOR_RGB2BGR))
                                byproducts.append(postmortem_file)
                        if has_bb:
                            start = i
                            should_analyze_file = False
                            current_zone = np.array(im[y0: y1, x0: x1])
                            current_file_path = f"{self.work_directory}/{prefix_splitted}-bounded-{i}.jpg"
                            cv2.imwrite(current_file_path, current_zone)
                        else:
                            current_zone = None
                            current_file_path = None
                            start = None
                        os.remove(resized_file_path)
                    os.remove(splitted_file)
                else:
                    if current_zone is not None:
                        data.extend([[current_file_path, start, i - 1]])
                    logger.info(f"--- inference done on {i} files")
                    break
                i += 1
                if i % 20 == 0:
                    logger.debug(f"------ {i} files parsed")
            annotations = pd.DataFrame(columns=['filename', 'start', 'end'], data=data)
            annotations.to_csv(annotations_file_path, encoding='utf-8')
        return annotations_file_path, byproducts

    @staticmethod
    def same_subs(expected_zone: np.ndarray, image_zone: np.ndarray) -> bool:
        threshold = 220
        expected_zone_thresholded = np.where(expected_zone > threshold, 1, 0)
        image_zone_thresholded = np.where(image_zone > threshold, 1, 0)
        ratio = (expected_zone_thresholded * image_zone_thresholded).sum() / (expected_zone_thresholded.sum() + 1e-6)
        return ratio > 0.75

    def add_subs(self, annotation_file_extraction: str, prefix: str) -> str:
        logger.info("Adding subs")
        file_with_subs_path = f"{self.work_directory}/{prefix}-with-subs.csv"
        if os.path.exists(file_with_subs_path):
            logger.info("--- Subs already added")
        else:
            df_annotations_in = pd.read_csv(annotation_file_extraction)
            background_images = []
            background_image = np.full((GOOGLE_MAX_HEIGHT, GOOGLE_MAX_WIDTH, 3), 255)
            row = 0
            column = 0
            has_data = False
            subtitles_per_frame = []
            nb_subs_for_page = 0
            for i, annotation in tqdm(df_annotations_in.iterrows()):
                y0 = row * COLUMN_HEIGHT
                x0 = column * COLUMN_WIDTH
                image_file_path = annotation['filename']
                coloured_image = read_image(image_file_path)
                height, width, _ = coloured_image.shape
                ydelta = min(height, COLUMN_HEIGHT)
                xdelta = min(width, COLUMN_WIDTH)
                y1 = y0 + ydelta
                x1 = x0 + xdelta
                # On part de en bas à droite et non en haut à gauche pour éviter les problèmes avec les images mal
                # découpées (en bas à droite on a plus de chances de trouver des sous-titres que en haut à gauche dans
                # une image mal découpée)
                background_image[y0:y1, x0: x1, :] = coloured_image[-ydelta:, -xdelta:, :]
                has_data = True
                column += 1
                nb_subs_for_page += 1
                if column >= NB_COLUMNS:
                    column = 0
                    row += 1
                    if row >= NB_ROWS:
                        background_images.append({'image': background_image, 'nb': nb_subs_for_page})
                        background_image = np.full((GOOGLE_MAX_HEIGHT, GOOGLE_MAX_WIDTH, 3), 255)
                        has_data = False
                        row = 0
                        nb_subs_for_page = 0
            if has_data:
                background_images.append({'image': background_image, 'nb': nb_subs_for_page})

            for i, bg in tqdm(enumerate(background_images)):
                bg_image_path = f"{self.work_directory}/{i}.jpg"
                cv2.imwrite(bg_image_path, bg['image'])
                full_text_annotation = self.send_image_to_google(bg_image_path)
                subs = self.get_texts(full_text_annotation)
                for sub in subs[: bg['nb']]:
                    subtitles_per_frame.append(self.__clean_sub___(sub))
            logger.info(f"{len(background_images)} calls made to Google")

            df_annotations_in['subs'] = subtitles_per_frame
            df_annotations_in = df_annotations_in[df_annotations_in.subs != '']
            df_annotations_in.to_csv(file_with_subs_path, encoding='utf-8')
        return file_with_subs_path

    def __clean_sub___(self, subs: str) -> str:
        subs = re.sub("\[.+?]", "", subs)
        subs = re.sub("\[.+?", "", subs)
        subs = re.sub("\(.+?\)", "", subs)
        subs = re.sub("\(.+?", "", subs)
        return subs.strip()

    def send_image_to_google(self, bg_image_path):
        logger.debug(f"........... Appel à Google")
        with io.open(bg_image_path, 'rb') as image_file:
            content = image_file.read()
        image = vision.Image(content=content)
        response = self.vision_client.document_text_detection(image=image)
        return response.full_text_annotation

    def get_block_coord(self, block) -> int:
        bounding_box = block.bounding_box
        first_vertex = bounding_box.vertices[0]
        x0, y0 = first_vertex.x, first_vertex.y
        row = int(round(y0 / COLUMN_HEIGHT))
        column = int(round(x0 / COLUMN_WIDTH))
        index = min(row * NB_COLUMNS + column, NB_COLUMNS * NB_ROWS - 1)
        return index

    def get_texts(self, document) -> str:
        my_texts = [''] * (NB_ROWS * NB_COLUMNS)
        for page in document.pages:
            for block in page.blocks:
                text = []
                for paragraph in block.paragraphs:
                    for word in paragraph.words:
                        for symbol in word.symbols:
                            if symbol.text:
                                if hasattr(symbol, 'property') and \
                                        (hasattr(symbol.property, 'detectedBreak') or
                                         hasattr(symbol.property, 'detected_break')):
                                    if hasattr(symbol.property, 'detectedBreak'):
                                        detected_break = symbol.property.detectedBreak.type
                                    else:
                                        detected_break = symbol.property.detected_break.type_
                                    if detected_break == BREAKS.UNKNOWN:
                                        break_symbol = ''
                                    elif detected_break == 'LINE_BREAK':
                                        break_symbol = '\n'
                                    else:
                                        break_symbol = ' '
                                else:
                                    break_symbol = ''
                                if symbol.text:
                                    text.append(f"{symbol.text}{break_symbol}")
                if len(text) > 0:
                    block_index = self.get_block_coord(block)
                    my_texts[block_index] = ''.join(text)
        return my_texts

    def _get_model(self, model_name: str) -> nn.Module:
        weights_path: str = f"/tmp/{model_name}"
        model = to_best_device(AttentionSubsBoxerModel(num_encoder_layers=8))
        self.storage.download(from_path=f"models/{model_name}", to_path=weights_path)
        loaded_weights = torch.load(weights_path, map_location=torch.device('cpu'))
        model.load_state_dict(loaded_weights)
        model.eval()
        return model

    def _is_trustable_bounding_box(self, x0: int, y0: int, x1: int, y1: int) -> bool:
        return x1 - x0 > self.bb_margin and y1 - y0 > self.bb_margin

    def clean(self, video_name):
        logger.info(f"Cleaning work files of {video_name}")
        filename_without_extension = os.path.splitext(os.path.basename(video_name))[0]
        file_list = glob.glob(f"{self.work_directory}/*{filename_without_extension}*")
        for file_path in tqdm(file_list):
            if os.path.isfile(file_path):
                try:
                    logger.debug(f"Removing file {file_path}")
                    os.remove(file_path)
                except BaseException as err:
                    logger.warning(f"Error while deleting file {file_path}", err)


from bs4 import BeautifulSoup
import requests
import re

url = 'https://keithgalli.github.io/web-scraping/webpage.html'

page = requests.get(url)

soup =BeautifulSoup(page.text, 'html.parser')

# look for all h1, h2, h3

headers = soup.find_all(['h1', 'h2', 'h3'])

#print(headers)

div = soup.find(class_='fun-facts')

# use regular expression
text = div.find(text=re.compile('dream'))
text = div.find_all(text=re.compile('d|D'))

# use select
paragraphs = soup.select("body > p")
print(paragraphs)

for paragraph in paragraphs:
  print(paragraph.select("i"))

soup.select("[align=middle]")

 #------------
files = page.select("div.block a")
relative_files = [f['href'] for f in files]

url = "https://keithgalli.github.io/web-scraping/"
for f in relative_files:
  full_url = url + f
  page = requests.get(full_url)
  bs_page = bs(page.content)
  secret_word_element = bs_page.find("p", attrs={"id": "secret-word"})
  secret_word = secret_word_element.string
  print(secret_word)


import os

from typing import List, Optional

from clipped.utils.paths import create_tarfile, get_files_in_path

from polyaxon import settings

def tar_dir(download_path: str) -> str:
    outputs_files = get_files_in_path(download_path)
    tar_base_name = os.path.basename(download_path)
    tar_name = "{}.tar.gz".format(tar_base_name)
    target_path = os.path.join(settings.CLIENT_CONFIG.archives_root or "", tar_name)
    create_tarfile(files=outputs_files, tar_path=target_path, relative_to=download_path)
    return target_path

def tar_files(
    download_name: str, outputs_files: List[str], relative_to: Optional[str] = None
) -> str:
    tar_name = "{}.tar.gz".format(download_name)
    target_path = os.path.join(settings.CLIENT_CONFIG.archives_root or "", tar_name)
    create_tarfile(files=outputs_files, tar_path=target_path, relative_to=relative_to)
    return target_path


#!/usr/bin/env python3

import argparse
import sqlite3
import pandas as pd
from datetime import datetime

# Parse command-line arguments
parser = argparse.ArgumentParser()
parser.add_argument(
    "--database_file",
    help="SQLite database file produced by NCBImeta.")
parser.add_argument(
    "--output",
    help="Fileto save (CSV-formatted) nextstrain metadata to.")
args = parser.parse_args()

# Connect to database
db_connection = sqlite3.connect(args.database_file)
# db_cursor = db_connection.cursor()

# Query database
query = "SELECT * FROM Assembly a LEFT JOIN BioSample b ON a.AssemblyBioSampleID = b.BioSample_id"
table = pd.read_sql_query(query, db_connection)
nextstrain_metadata = {
    "strain": table['AssemblyAccession'],
    "date": table['BioSampleCollectionDate'],
    "date_submitted": table['AssemblySubmissionDate'],
    "region": table['BioSampleGeographicLocation'],
    "host": table['BioSampleHost']
}

# Parse dates, function inspired by https://github.com/nextstrain/ncov-ingest/tree/master
def format_date(date_string: str) -> str:
    if date_string is None:
        return '?'
    try:
        return datetime.strptime(date_string, '%Y-%m-%d').strftime('%Y-%m-%d')
    except ValueError:
        try:
            return datetime.strptime(date_string, '%Y-%m').strftime('%Y-%m-XX')
        except ValueError:
            try:
                return datetime.strptime(date_string, '%Y').strftime('%Y-XX-XX')
            except:
                return '?'

# Format nicely
for column in ['date', 'date_submitted']:
    if column in nextstrain_metadata:
        new_dates = [None] * len(nextstrain_metadata[column])
        for i, date in enumerate(nextstrain_metadata[column]):
            new_date = format_date(date)
            new_dates[i] = new_date
        nextstrain_metadata[column] = new_dates
nextstrain_table = pd.DataFrame(nextstrain_metadata)
nextstrain_table.fillna('?', inplace=True)

# Write out nextstrain metadata
nextstrain_table.to_csv(args.output, index=False)


import imageio
import numpy as np

image = np.array(imageio.imread('output1.exr'))
image = image ** 2.2
brightness = 0.2126 * image[..., 0] + \
             0.7152 * image[..., 1] + \
             0.0722 * image[..., 2]
image[..., 0] /= brightness
image[..., 1] /= brightness
image[..., 2] /= brightness
print('mean %s' % np.mean(brightness))
print(np.mean(brightness) + 3*(np.var(brightness)**0.5))
maxbr = 5
factor = 1000
brightness = np.log(brightness * factor + 1) / np.log(maxbr * factor + 1)
w, h, dims = image.shape
result = np.zeros((w, h, 4))
result[..., 0:3] = (image ** (1 / 2.2)) * 255
result[..., 3] = brightness * 255

result = np.rint(np.minimum(result, 255.0))

imageio.imwrite('output1.png', result)


import inspect

import ctk
import slicer
import numpy
import time
import qt
import vtk
import colorsys
from slicer.ScriptedLoadableModule import *

# Interface tools
class InterfaceTools:
    def __init__(self):
        pass

    @staticmethod
    def build_vertical_space(height=18):
        s = qt.QWidget()
        s.setFixedSize(10, height)
        return s

    @staticmethod
    def build_dropdown(title, disabled=False):
        d = ctk.ctkCollapsibleButton()
        d.text = title
        d.enabled = not disabled
        d.collapsed = disabled
        return d

    @staticmethod
    def build_frame():
        g = qt.QFrame()
        g.setFrameStyle(3)
        return g

    @staticmethod
    def build_volume_selector(on_click=None):
        s = slicer.qMRMLNodeComboBox()
        s.nodeTypes = ["vtkMRMLScalarVolumeNode"]
        s.noneEnabled = True
        s.addEnabled = False
        s.renameEnabled = True
        s.setMRMLScene(slicer.mrmlScene)
        if on_click is not None: s.connect("currentNodeChanged(bool)", on_click)
        return s

    @staticmethod
    def build_model_selector(on_click=None):
        s = slicer.qMRMLNodeComboBox()
        s.nodeTypes = ["vtkMRMLModelNode"]
        s.addEnabled = False
        s.renameEnabled = s.noneEnabled = False  # True
        s.setMRMLScene(slicer.mrmlScene)
        if on_click is not None: s.connect("currentNodeChanged(bool)", on_click)
        return s

    @staticmethod
    def build_combo_box(items, current_index_changed):
        comboBox = qt.QComboBox()
        comboBox.addItems(items)
        comboBox.setFixedWidth(350)
        comboBox.connect("currentIndexChanged(QString)", current_index_changed)
        box = qt.QHBoxLayout()
        box.addStretch()
        box.addWidget(comboBox)
        return box

    @staticmethod
    def build_spin_box(minimum, maximum, click=None, decimals=0, step=1.0, initial=0.0, width=None):
        box = qt.QDoubleSpinBox()
        box.setMinimum(minimum)
        box.setMaximum(maximum)
        box.setSingleStep(step)
        box.setDecimals(decimals)
        box.setValue(initial)
        if width is not None: box.setFixedWidth(width)
        if click is not None: box.connect('valueChanged(double)', click)
        return box

    @staticmethod
    def build_radio_button(title, on_click, checked=False, tooltip=None, width=None):
        b = qt.QRadioButton(title)
        b.connect('clicked(bool)', on_click)
        b.setChecked(checked)
        if tooltip is not None: b.setToolTip(tooltip)
        if width is not None: b.setFixedWidth(width)
        return b

    @staticmethod
    def build_label(text, width=None):
        b = qt.QLabel(text)
        if width is not None: b.setFixedWidth(width)
        return b

    @staticmethod
    def build_icon_button(icon_path, on_click, width=50, tooltip=None):
        path = slicer.os.path.dirname(slicer.os.path.abspath(inspect.getfile(inspect.currentframe()))) + icon_path
        icon = qt.QPixmap(path).scaled(qt.QSize(16, 16), qt.Qt.KeepAspectRatio, qt.Qt.SmoothTransformation)
        b = qt.QToolButton()
        b.setIcon(qt.QIcon(icon))
        b.setFixedSize(width, 24)
        if tooltip is not None: b.setToolTip(tooltip)
        b.connect('clicked(bool)', on_click)
        return b

    @staticmethod
    def build_min_max(initial, decimals=2, step=0.1, lb=0.0, hb=1000.0, units='mm', min_text='MIN: ', max_text='MAX: '):
        def set_min(value): initial[0] = value
        def set_max(value): initial[1] = value
        box = qt.QHBoxLayout()
        box.addStretch()
        box.addWidget(InterfaceTools.build_label(min_text, 40))
        lowerSb = InterfaceTools.build_spin_box(lb, hb, decimals=decimals, step=step, initial=initial[0], click=set_min, width=80)
        box.addWidget(lowerSb)
        box.addWidget(InterfaceTools.build_label(units, 50))
        box.addWidget(InterfaceTools.build_label(max_text, 40))
        higherSb = InterfaceTools.build_spin_box(lb, hb, decimals=decimals, step=step, initial=initial[1], click=set_max, width=80)
        box.addWidget(higherSb)
        box.addWidget(InterfaceTools.build_label(units, 30))

        def set_boxes(value=None, enabled=None):
            if value is not None:
                set_min(value[0])
                lowerSb.setValue(value[0])
                set_max(value[1])
                higherSb.setValue(value[1])
            if enabled is not None:
                lowerSb.enabled = enabled
                higherSb.enabled = enabled
        return box, set_boxes

class BoneThicknessMappingType:
    THICKNESS = 'Thickness to dura'
    AIR_CELL = 'Distance to first air cell'

class BoneSegmentationLowerBound:
    MANUAL = 'Manual'
    CLINICAL = 'Clinical CT'
    CONE_BEAM = 'Cone Beam CT'

class BoneDepthMappingPresets:
    MANUAL = 'Manual'
    BCI601 = 'BCI 601'
    BCI602 = 'BCI 602'

class BoneThicknessMappingState:
    WAITING = 1
    READY = 2
    EXECUTING = 3
    FINISHED = 4

class BoneThicknessMappingQuality:
    VERY_LOW = 'VERY LOW (ray every 4 dimensional units)'
    LOW = 'LOW (ray every 2 dimensional units)'
    MEDIUM = 'MEDIUM (ray every 1 dimensional unit)'
    HIGH = 'HIGH (ray every 0.5 dimensional units)'
    VERY_HIGH = 'VERY HIGH (ray every 0.25 dimensional units)'
    EXTREME = 'EXTREME (ray every 0.125 dimensional units)'

class HitPoint:
    pid = None
    point = None
    normal = [0.0, 0.0, 0.0]

    def __init__(self, pid, point):
        self.pid = pid
        self.point = point

class BoneThicknessMapping(ScriptedLoadableModule):
    def __init__(self, parent):
        ScriptedLoadableModule.__init__(self, parent)
        self.parent.title = "Bone Thickness Mapping"
        self.parent.categories = ["Shape Analysis"]
        self.parent.dependencies = []
        self.parent.contributors = ["Evan Simpson (Western University)"]
        self.parent.helpText = "The following module will segment and threshold a volume to isolate bone material, cast rays in one direction to calculate the thickness of segmented bone, finally a gradient visualization will be rendered on the 3D segment model." \
                               "\n[Version 1.2-2020.06.18]"
        self.parent.acknowledgementText = "This module was originally developed by Evan Simpson at The University of Western Ontario in the HML/SKA Auditory Biophysics Lab."

class BoneThicknessMappingWidget(ScriptedLoadableModuleWidget):
    # Data members --------------
    state = BoneThicknessMappingState.WAITING
    status, progress = 'N/A', 0
    thicknessScalarArray, airCellScalarArray = None, None
    thicknessColourNode, airCellColourNode = None, None
    modelPolyData = None
    segmentationBounds = None
    topLayerPolyData = None
    hitPointList = None
    modelNode = None

    # Configuration preferences
    CONFIG_precision = 1.0
    CONFIG_rayCastAxis = ctk.ctkAxesWidget.Left
    CONFIG_segmentThresholdRange = [600, 3071]
    CONFIG_regionOfInterest = [-100, 100]
    CONFIG_minMaxAirCell = [0.0, 4.0]
    CONFIG_minMaxSkullThickness = [0.0, 8.7]
    CONFIG_mmOfAirPastBone = 4.0

    # UI members (in order of appearance) --------------
    infoLabel = None
    volumeSelector = None
    configuration_tools = None
    statusLabel = None
    executeButton = None
    progressBar = None
    finishButton = None
    resultSection = None
    displayThicknessSelector = None
    displayFirstAirCellSelector = None
    displayScalarBarCheckbox = None

    def __init__(self, parent=None):
        ScriptedLoadableModuleWidget.__init__(self, parent)

    def setup(self):
        ScriptedLoadableModuleWidget.setup(self)
        self.layout.addLayout(self.build_input_tools())
        self.layout.addLayout(self.build_execution_tools())
        self.layout.addLayout(self.build_result_tools())
        self.layout.addStretch()
        BoneThicknessMappingLogic.reset_view(ctk.ctkAxesWidget.Left)
        self.update_all()
        slicer.app.aboutToQuit.connect(self.release_memory)

    # interface build ------------------------------------------------------------------------------
    def build_input_tools(self):
        layout = qt.QVBoxLayout()

        self.volumeSelector = InterfaceTools.build_volume_selector(on_click=self.click_input_selector)
        box = qt.QHBoxLayout()
        box.addWidget(self.volumeSelector)
        box.addWidget(InterfaceTools.build_icon_button('/Resources/Icons/fit.png', on_click=lambda: BoneThicknessMappingLogic.reset_view(self.CONFIG_rayCastAxis), tooltip="Reset 3D view."))
        form = qt.QFormLayout()
        form.addRow(qt.QLabel('Select an input volume to auto-segment, render, and calculate thickness.'))
        form.addRow("Input Volume: ", box)
        self.infoLabel = qt.QLabel()
        form.addWidget(self.infoLabel)

        title = qt.QHBoxLayout()
        path = slicer.os.path.dirname(slicer.os.path.abspath(inspect.getfile(inspect.currentframe()))) + "/Resources/Icons/logo.png"
        icon = qt.QPixmap(path).scaled(qt.QSize(240, 400), qt.Qt.KeepAspectRatio, qt.Qt.SmoothTransformation)
        logo = qt.QLabel()
        logo.setPixmap(icon)
        title.addWidget(logo)

        layout.addLayout(title)
        layout.addLayout(form)
        layout.setMargin(10)
        return layout

    def build_configuration_tools(self):
        self.configuration_tools = InterfaceTools.build_dropdown("Configuration")
        layout = qt.QFormLayout(self.configuration_tools)

        # threshold
        threshBox, setThresh = InterfaceTools.build_min_max(self.CONFIG_segmentThresholdRange, step=5.0, decimals=2, lb=-4000, hb=4000, units='')

        def current_index_changed(string):
            if string == BoneSegmentationLowerBound.MANUAL:
                setThresh([600.0, 3071.0], enabled=True)
            elif string == BoneSegmentationLowerBound.CLINICAL:
                setThresh([650.0, 3071.0], enabled=False)
            elif string == BoneSegmentationLowerBound.CONE_BEAM:
                setThresh([750.0, 3071.0], enabled=False)
        presets = InterfaceTools.build_combo_box(
            items=[BoneSegmentationLowerBound.MANUAL, BoneSegmentationLowerBound.CLINICAL, BoneSegmentationLowerBound.CONE_BEAM],
            current_index_changed=current_index_changed
        )

        group_box = qt.QGroupBox('Auto segmenting')
        group_layout = qt.QFormLayout(group_box)
        group_layout.addRow("Presets", presets)
        group_layout.addRow("Otsu bone-threshold range", threshBox)
        layout.addRow(group_box)

        # ray direction
        def set_axis(a):
            self.CONFIG_rayCastAxis = a
            slicer.app.layoutManager().threeDWidget(0).threeDView().lookFromViewAxis(a)

        dirBox = qt.QVBoxLayout()
        row1 = qt.QHBoxLayout()
        row1.addStretch()
        row1.addWidget(InterfaceTools.build_radio_button('R', lambda: set_axis(ctk.ctkAxesWidget.Right), width=100))
        row1.addWidget(InterfaceTools.build_radio_button('A', lambda: set_axis(ctk.ctkAxesWidget.Anterior), width=100))
        row1.addWidget(InterfaceTools.build_radio_button('S', lambda: set_axis(ctk.ctkAxesWidget.Superior), width=100))
        row2 = qt.QHBoxLayout()
        row2.addStretch()
        row2.addWidget(InterfaceTools.build_radio_button('L', lambda: set_axis(ctk.ctkAxesWidget.Left), width=100, checked=True))
        row2.addWidget(InterfaceTools.build_radio_button('P', lambda: set_axis(ctk.ctkAxesWidget.Posterior), width=100))
        row2.addWidget(InterfaceTools.build_radio_button('I', lambda: set_axis(ctk.ctkAxesWidget.Inferior), width=100))
        dirBox.addLayout(row1)
        dirBox.addLayout(row2)

        # mm of air afterbone
        airBox = qt.QHBoxLayout()
        def set_air(mm): self.CONFIG_mmOfAirPastBone = mm
        airBox.addStretch()
        airBox.addWidget(InterfaceTools.build_spin_box(0.0, 1000.0, decimals=2, click=set_air, step=0.1, initial=self.CONFIG_mmOfAirPastBone, width=260))
        airBox.addWidget(InterfaceTools.build_label("mm", width=30))

        # region of interest
        roiBox, setRoi = InterfaceTools.build_min_max(self.CONFIG_regionOfInterest, step=1.0, decimals=0, lb=-1000, hb=1000, units='units')

        # add ray-casting box
        group_box = qt.QGroupBox('Ray-casting')
        group_layout = qt.QFormLayout(group_box)
        group_layout.addRow("Cast direction: ", dirBox)
        group_layout.addRow("Air allowance after bone: ", airBox)
        group_layout.addRow("Casting bounds (along cast direction):", roiBox)
        layout.addRow(InterfaceTools.build_vertical_space())
        layout.addRow(group_box)

        # degrees of interest
        # TODO

        # min/max
        skullBox, setSkullBoxes = InterfaceTools.build_min_max(self.CONFIG_minMaxSkullThickness)
        airCellBox, setAirBoxes = InterfaceTools.build_min_max(self.CONFIG_minMaxAirCell)

        def pick_depth_preset(string):
            if string == BoneDepthMappingPresets.MANUAL:
                setSkullBoxes(enabled=True)
                setAirBoxes(enabled=True)
            elif string == BoneDepthMappingPresets.BCI601:
                setSkullBoxes([0.0, 8.7], enabled=False)
                setAirBoxes([0.0, 4.0], enabled=False)
            elif string == BoneDepthMappingPresets.BCI602:
                setSkullBoxes([0.0, 4.5], enabled=False)
                setAirBoxes([0.0, 4.0], enabled=False)
        comboBox = InterfaceTools.build_combo_box(
            items=[BoneDepthMappingPresets.MANUAL, BoneDepthMappingPresets.BCI601, BoneDepthMappingPresets.BCI602],
            current_index_changed=pick_depth_preset
        )

        # add ray-casting box
        group_box = qt.QGroupBox('Depth mapping')
        g_layout = qt.QFormLayout(group_box)
        g_layout.addRow('Depth preset: ', comboBox)
        g_layout.addRow("Thickness depth: ", skullBox)
        g_layout.addRow("Air cell depth: ", airCellBox)
        layout.addRow(InterfaceTools.build_vertical_space())
        layout.addRow(group_box)

        # quality
        def current_index_changed(string):
            if string == BoneThicknessMappingQuality.VERY_LOW: self.CONFIG_precision = 4.0
            elif string == BoneThicknessMappingQuality.LOW: self.CONFIG_precision = 2.0
            elif string == BoneThicknessMappingQuality.MEDIUM: self.CONFIG_precision = 1.0
            elif string == BoneThicknessMappingQuality.HIGH: self.CONFIG_precision = 0.50
            elif string == BoneThicknessMappingQuality.VERY_HIGH: self.CONFIG_precision = 0.25
            elif string == BoneThicknessMappingQuality.EXTREME: self.CONFIG_precision = 0.125
        comboBox = qt.QComboBox()
        comboBox.addItems([BoneThicknessMappingQuality.VERY_LOW, BoneThicknessMappingQuality.LOW, BoneThicknessMappingQuality.MEDIUM, BoneThicknessMappingQuality.HIGH, BoneThicknessMappingQuality.VERY_HIGH, BoneThicknessMappingQuality.EXTREME])
        comboBox.setCurrentIndex(2)
        comboBox.setFixedWidth(350)
        comboBox.connect("currentIndexChanged(QString)", current_index_changed)
        box = qt.QHBoxLayout()
        box.addStretch()
        box.addWidget(comboBox)

        # add ray-casting box
        group_box = qt.QGroupBox('Rendering')
        g_layout = qt.QFormLayout(group_box)
        g_layout.addRow("Render quality: ", box)
        layout.addRow(InterfaceTools.build_vertical_space())
        layout.addRow(group_box)

        layout.addRow(InterfaceTools.build_vertical_space())
        layout.setMargin(10)
        return self.configuration_tools

    def build_execution_tools(self):
        self.executeButton = qt.QPushButton('Execute')
        self.executeButton.setFixedHeight(36)
        self.executeButton.connect('clicked(bool)', self.click_execute)
        self.statusLabel = qt.QLabel('Status: ')
        self.statusLabel.enabled = False
        self.progressBar = qt.QProgressBar()
        self.progressBar.setFixedHeight(36)
        self.progressBar.minimum = 0
        self.progressBar.maximum = 100
        self.progressBar.value = 0
        self.progressBar.visible = False
        self.finishButton = qt.QPushButton('Finish')
        self.finishButton.visible = False
        self.finishButton.connect('clicked(bool)', self.click_finish)
        self.finishButton.setFixedHeight(36)
        box = qt.QHBoxLayout()
        box.addWidget(self.progressBar)
        box.addWidget(self.finishButton)

        layout = qt.QVBoxLayout()
        layout.addWidget(self.build_configuration_tools())
        layout.addWidget(self.statusLabel)
        layout.addWidget(self.executeButton)
        layout.addLayout(box)
        layout.setMargin(10)
        return layout

    def build_result_tools(self):
        self.resultSection = InterfaceTools.build_frame()

        self.displayThicknessSelector = InterfaceTools.build_radio_button(BoneThicknessMappingType.THICKNESS, self.click_result_radio, checked=True)
        self.displayFirstAirCellSelector = InterfaceTools.build_radio_button(BoneThicknessMappingType.AIR_CELL, self.click_result_radio)

        box = qt.QVBoxLayout()
        box.addWidget(self.displayThicknessSelector)
        box.addWidget(self.displayFirstAirCellSelector)

        self.displayScalarBarCheckbox = qt.QCheckBox()
        self.displayScalarBarCheckbox.checked = True
        self.displayScalarBarCheckbox.connect("stateChanged(int)", self.click_toggle_scalar_bar)

        form = qt.QFormLayout(self.resultSection)
        # form.addRow('Results', qt.QWidget())
        form.addRow("Map Display: ", box)
        # form.addRow(qt.QLayout())
        form.addRow("Display Scalar Bar: ", self.displayScalarBarCheckbox)
        form.setContentsMargins(10, 8, 10, 14)

        layout = qt.QVBoxLayout()
        layout.addWidget(self.resultSection)
        layout.setContentsMargins(12, 0, 12, 10)
        return layout

    # interface update ------------------------------------------------------------------------------
    def update_all(self):
        self.update_input_tools()
        self.update_execution_tools()
        self.update_results()

    def update_input_tools(self):
        if self.volumeSelector.currentNode() is None:
            self.volumeSelector.enabled = True
            self.infoLabel.text = 'Dimensions: '
            self.infoLabel.enabled = False
            self.state = BoneThicknessMappingState.WAITING
        elif self.state is BoneThicknessMappingState.WAITING and self.volumeSelector.currentNode() is not None:
            self.volumeSelector.enabled = True
            self.infoLabel.enabled = True
            self.infoLabel.text = 'Dimensions: ' + str(self.volumeSelector.currentNode().GetImageData().GetDimensions())
            self.state = BoneThicknessMappingState.READY
            # self.modelPolyData = SkullThicknessMappingLogic.process_segmentation(image=self.volumeSelector.currentNode(), update_status=self.update_status)
        elif self.state is BoneThicknessMappingState.EXECUTING or self.state is BoneThicknessMappingState.FINISHED:
            self.volumeSelector.enabled = False

    def update_execution_tools(self):
        if self.state is BoneThicknessMappingState.WAITING:
            self.configuration_tools.enabled = True
            self.configuration_tools.collapsed = True
            self.executeButton.visible = True
            self.executeButton.enabled = False
            self.progressBar.visible = False
            self.progressBar.value = 0
            self.statusLabel.enabled = False
            self.statusLabel.text = 'Status: WAITING'
            self.finishButton.visible = False
        elif self.state is BoneThicknessMappingState.READY:
            self.configuration_tools.enabled = True
            self.executeButton.visible = True
            self.executeButton.enabled = True
            self.progressBar.visible = False
            self.progressBar.value = 0
            self.statusLabel.enabled = True
            self.statusLabel.text = 'Status: READY'
            self.finishButton.visible = False
        elif self.state is BoneThicknessMappingState.EXECUTING:
            self.configuration_tools.enabled = False
            self.configuration_tools.collapsed = True
            self.executeButton.visible = False
            self.executeButton.enabled = False
            self.progressBar.visible = True
            self.progressBar.value = self.progress
            self.statusLabel.enabled = True
            self.statusLabel.text = 'Status: ' + str(self.status)
            self.finishButton.visible = False
        elif self.state is BoneThicknessMappingState.FINISHED:
            self.configuration_tools.enabled = False
            self.configuration_tools.collapsed = True
            self.progressBar.value = 100
            self.executeButton.visible = False
            self.finishButton.visible = True

    def update_results(self):
        if self.thicknessScalarArray is not None and self.airCellScalarArray is not None:
            self.resultSection.enabled = True
        else: self.resultSection.enabled = False

    def update_status(self, text=None, progress=None):
        if text is not None:
            print(text)
            self.status = str(text)
        if progress is not None:
            self.progress = progress
        self.update_all()
        slicer.app.processEvents()

    # interface click events ----------------------------------------------------------------------
    def click_input_selector(self):
        if self.volumeSelector.currentNode() is not None:
            BoneThicknessMappingLogic.update_input_volume(self.volumeSelector.currentNode().GetID())
        self.update_all()

    def click_execute(self):
        # TODO add try and catch
        if self.state is not BoneThicknessMappingState.READY: return
        self.state = BoneThicknessMappingState.EXECUTING
        self.update_status(text='Initializing execution..', progress=0)
        BoneThicknessMappingLogic.reset_view(self.CONFIG_rayCastAxis)
        BoneThicknessMappingLogic.clear_3d_view()
        BoneThicknessMappingLogic.set_scalar_colour_bar_state(0)
        self.modelPolyData, self.segmentationBounds = BoneThicknessMappingLogic.process_segmentation(
            threshold_range=self.CONFIG_segmentThresholdRange,
            image=self.volumeSelector.currentNode(),
            axis=self.CONFIG_rayCastAxis,
            update_status=self.update_status
        )
        self.topLayerPolyData, self.hitPointList = BoneThicknessMappingLogic.rainfall_quad_cast(
            poly_data=self.modelPolyData,
            seg_bounds=self.segmentationBounds,
            cast_axis=self.CONFIG_rayCastAxis,
            precision=self.CONFIG_precision,
            region_of_interest=self.CONFIG_regionOfInterest,
            update_status=self.update_status
        )
        self.modelNode = BoneThicknessMappingLogic.build_model(
            poly_data=self.topLayerPolyData,
            update_status=self.update_status
        )
        self.thicknessScalarArray, self.airCellScalarArray = BoneThicknessMappingLogic.ray_cast_color_thickness(
            poly_data=self.modelPolyData,
            hit_point_list=self.hitPointList,
            cast_axis=self.CONFIG_rayCastAxis,
            dimensions=self.volumeSelector.currentNode().GetImageData().GetDimensions(),
            mm_of_air_past_bone=self.CONFIG_mmOfAirPastBone,
            update_status=self.update_status
        )
        self.thicknessColourNode, self.airCellColourNode = BoneThicknessMappingLogic.build_color_table_nodes(
            minmax_thickness=self.CONFIG_minMaxSkullThickness,
            minmax_air_cell=self.CONFIG_minMaxAirCell
        )
        # finalize
        self.click_result_radio()
        self.state = BoneThicknessMappingState.FINISHED
        self.update_status(progress=100)

    def click_finish(self):
        self.state = BoneThicknessMappingState.WAITING
        self.update_all()

    def click_result_radio(self):
        if self.thicknessScalarArray is None or self.airCellScalarArray is None: return  # TODO add error message
        scalar, scalarName, colourNodeId = None, None, None
        if self.displayThicknessSelector.isChecked():
            scalar = self.thicknessScalarArray
            scalarName = BoneThicknessMappingType.THICKNESS
            colourNodeId = self.thicknessColourNode.GetID()
        elif self.displayFirstAirCellSelector.isChecked():
            scalar = self.airCellScalarArray
            scalarName = BoneThicknessMappingType.AIR_CELL
            colourNodeId = self.airCellColourNode.GetID()
        # update poly data
        self.topLayerPolyData.GetPointData().SetScalars(scalar)
        self.topLayerPolyData.Modified()
        # update display node
        displayNode = self.modelNode.GetDisplayNode()
        displayNode.SetActiveScalarName(scalarName)
        displayNode.SetAndObserveColorNodeID(colourNodeId)
        displayNode.ScalarVisibilityOn()
        displayNode.SetScalarRangeFlag(slicer.vtkMRMLDisplayNode.UseColorNodeScalarRange)
        # update scalar bar
        BoneThicknessMappingLogic.set_scalar_colour_bar_state(1, colourNodeId)
        # reset view
        # BoneThicknessMappingLogic.reset_view(self.CONFIG_rayCastAxis)

    def click_toggle_scalar_bar(self, state):
        if state == 0: BoneThicknessMappingLogic.set_scalar_colour_bar_state(0)
        elif state == 2:
            if self.displayThicknessSelector.isChecked(): BoneThicknessMappingLogic.set_scalar_colour_bar_state(1, self.thicknessColourNode)
            elif self.displayFirstAirCellSelector.isChecked(): BoneThicknessMappingLogic.set_scalar_colour_bar_state(1, self.airCellColourNode)

    def release_memory(self):

        # UI
        self.infoLabel = None
        self.volumeSelector = None
        self.configuration_tools = None
        self.statusLabel = None
        self.executeButton = None
        self.progressBar = None
        self.finishButton = None
        self.resultSection = None
        self.displayThicknessSelector = None
        self.displayFirstAirCellSelector = None
        self.displayScalarBarCheckbox = None

        # Config
        self.CONFIG_precision = None
        self.CONFIG_rayCastAxis = None
        self.CONFIG_segmentThresholdRange = None
        self.CONFIG_regionOfInterest = None
        self.CONFIG_minMaxAirCell = None
        self.CONFIG_minMaxSkullThickness = None
        self.CONFIG_mmOfAirPastBone = None

        # Data
        self.thicknessScalarArray, self.airCellScalarArray = None, None
        self.thicknessColourNode, self.airCellColourNode = None, None
        self.modelPolyData = None
        self.segmentationBounds = None
        self.topLayerPolyData = None
        self.hitPointList = None
        self.modelNode = None

class BoneThicknessMappingLogic(ScriptedLoadableModuleLogic):
    @staticmethod
    def update_input_volume(volume_id):
        for c in ['Red', 'Yellow', 'Green']:
            slicer.app.layoutManager().sliceWidget(c).sliceLogic().GetSliceCompositeNode().SetBackgroundVolumeID(volume_id)
        logic = slicer.app.layoutManager().mrmlSliceLogics()
        for i in range(logic.GetNumberOfItems()): logic.GetItemAsObject(i).FitSliceToAll()

    @staticmethod
    def reset_view(axis):
        m = slicer.app.layoutManager()
        m.setLayout(16)
        w = m.threeDWidget(0)
        w.threeDView().lookFromViewAxis(axis)
        w.threeDView().renderWindow().GetRenderers().GetFirstRenderer().ResetCamera()
        # c = w.threeDController()
        # for i in range(12): c.zoomIn()

    @staticmethod
    def clear_3d_view():
        for n in slicer.mrmlScene.GetNodesByClass("vtkMRMLModelNode"): n.GetDisplayNode().SetVisibility(0)
        for n in slicer.mrmlScene.GetNodesByClass("vtkMRMLSegmentationNode"): n.GetDisplayNode().SetVisibility(0)
        # hide 3D cube & labels
        v = slicer.util.getNode("View1")
        v.SetBoxVisible(False)
        v.SetAxisLabelsVisible(False)

    @staticmethod
    def process_segmentation(threshold_range, image, axis, update_status):
        # Fix Volume Orientation
        update_status(text="Rotating views to volume plane...", progress=2)
        manager = slicer.app.layoutManager()
        for name in manager.sliceViewNames():
            widget = manager.sliceWidget(name)
            node = widget.mrmlSliceNode()
            node.RotateToVolumePlane(image)

        # Create segmentation
        update_status(text="Creating segmentation...", progress=5)
        segmentationNode = slicer.mrmlScene.AddNewNodeByClass("vtkMRMLSegmentationNode")
        segmentationNode.CreateDefaultDisplayNodes()
        segmentationNode.SetReferenceImageGeometryParameterFromVolumeNode(image)
        segmentId = segmentationNode.GetSegmentation().AddEmptySegment("Bone")
        segmentationNode.GetSegmentation().GetSegment(segmentId).SetColor([0.9, 0.8, 0.7])

        # Create segment editor to get access to effects
        update_status(text="Starting segmentation editor...", progress=6)
        segmentEditorWidget = slicer.qMRMLSegmentEditorWidget()
        segmentEditorWidget.setMRMLScene(slicer.mrmlScene)
        segmentEditorNode = slicer.mrmlScene.AddNewNodeByClass("vtkMRMLSegmentEditorNode")
        segmentEditorWidget.setMRMLSegmentEditorNode(segmentEditorNode)
        segmentEditorWidget.setSegmentationNode(segmentationNode)
        segmentEditorWidget.setMasterVolumeNode(image)

        # Threshold
        update_status(text="Processing threshold segmentation...", progress=8)
        segmentEditorWidget.setActiveEffectByName("Threshold")
        effect = segmentEditorWidget.activeEffect()
        effect.setParameter("MinimumThreshold", str(threshold_range[0]))  # 1460 #1160 # 223
        effect.setParameter("MaximumThreshold", str(threshold_range[1]))
        effect.self().onApply()

        # Smoothing
        update_status(text="Processing smoothing segmentation...", progress=10)
        segmentEditorWidget.setActiveEffectByName("Smoothing")
        effect = segmentEditorWidget.activeEffect()
        effect.setParameter("SmoothingMethod", "MORPHOLOGICAL_OPENING")
        effect.setParameter("KernelSizeMm", 0.5)
        effect.self().onApply()

        # Islands
        update_status(text="Processing island segmentation...", progress=11)
        segmentEditorWidget.setActiveEffectByName("Islands")
        effect = segmentEditorWidget.activeEffect()
        effect.setParameter("Operation", "KEEP_LARGEST_ISLAND")
        effect.setParameter("MinimumSize", 1000)
        effect.self().onApply()

        # Crop
        # update_status("Cropping segmentation...")
        # segmentEditorWidget.setActiveEffectByName("Scissors")
        # effect = segmentEditorWidget.activeEffect()
        # effect.setParameter("MinimumThreshold", "223")
        # effect.setParameter("MaximumThreshold", "3071")
        # effect.self().onApply()

        # Clean up
        update_status(text="Cleaning up...", progress=13)
        segmentEditorWidget.setActiveEffectByName(None)
        slicer.mrmlScene.RemoveNode(segmentEditorNode)

        # Make segmentation results visible in 3D and set focal
        update_status(text="Rendering...", progress=15)
        segmentationNode.CreateClosedSurfaceRepresentation()
        BoneThicknessMappingLogic.reset_view(axis)

        # Retrieve segmentation bounds
        bounds = [0, 0, 0, 0, 0, 0]
        segmentationNode.GetBounds(bounds)

        # Make sure surface mesh cells are consistently oriented
        update_status(text="Retrieving surface mesh...", progress=18)
        if slicer.app.majorVersion == 4 and slicer.app.minorVersion <= 10:
            polyData = segmentationNode.GetClosedSurfaceRepresentation(segmentId)
        else:
            polyData = vtk.vtkPolyData()
            segmentationNode.GetClosedSurfaceRepresentation(segmentId, polyData)

        return polyData, bounds

    @staticmethod
    def determine_cast_axis_index(cast_axis):
        castIndex = None
        if cast_axis in [ctk.ctkAxesWidget.Right, ctk.ctkAxesWidget.Left]: castIndex = 0
        elif cast_axis in [ctk.ctkAxesWidget.Anterior, ctk.ctkAxesWidget.Posterior]: castIndex = 1
        elif cast_axis in [ctk.ctkAxesWidget.Superior, ctk.ctkAxesWidget.Inferior]: castIndex = 2
        return castIndex

    @staticmethod
    def rainfall_quad_cast(poly_data, seg_bounds, cast_axis, precision, region_of_interest, update_status):
        update_status(text="Building intersection object tree...", progress=41)
        bspTree = vtk.vtkModifiedBSPTree()
        bspTree.SetDataSet(poly_data)
        bspTree.BuildLocator()

        update_status(text="Calculating segmentation dimensions...", progress=42)
        negated = 1 if cast_axis in [ctk.ctkAxesWidget.Right, ctk.ctkAxesWidget.Anterior, ctk.ctkAxesWidget.Superior] else -1
        castIndex = BoneThicknessMappingLogic.determine_cast_axis_index(cast_axis)
        castVector = [0.0, 0.0, 0.0]
        castVector[castIndex] = 1.0 * negated
        castPlaneIndices = [0, 1, 2]
        castPlaneIndices.remove(castIndex)

        update_status(text="Calculating segmentation cast-plane...", progress=43)
        depthIncrements = [seg_bounds[castIndex], seg_bounds[castIndex+1]]
        if cast_axis in [ctk.ctkAxesWidget.Right, ctk.ctkAxesWidget.Posterior, ctk.ctkAxesWidget.Superior]: depthIncrements.reverse()
        horizontalIncrements = [seg_bounds[castPlaneIndices[0]*2], seg_bounds[castPlaneIndices[0]*2+1]]
        verticalIncrements = [seg_bounds[castPlaneIndices[1]*2],  seg_bounds[castPlaneIndices[1]*2+1]]
        castPlaneIncrements = [
            int(abs(horizontalIncrements[0] - horizontalIncrements[1])/precision),
            int(abs(verticalIncrements[0] - verticalIncrements[1])/precision),
        ]

        def build_ray(i, j):
            p1 = [None, None, None]
            p1[castIndex] = depthIncrements[0] + negated*100
            p1[castPlaneIndices[0]] = horizontalIncrements[0] + i*precision
            p1[castPlaneIndices[1]] = verticalIncrements[0] + j*precision
            p2 = p1[:]
            p2[castIndex] = depthIncrements[1] - negated*100
            return p1, p2

        # cast rays
        update_status(text="Casting " + str(int(castPlaneIncrements[0]*castPlaneIncrements[1])) + " rays...", progress=44)
        startTime = time.time()
        points, temporaryHitPoint = vtk.vtkPoints(), [0.0, 0.0, 0.0]
        hitPointMatrix = [[None for j in range(castPlaneIncrements[1])] for i in range(castPlaneIncrements[0])]
        for i in range(len(hitPointMatrix)):
            for j in range(len(hitPointMatrix[i])):
                start, end = build_ray(i, j)
                res = bspTree.IntersectWithLine(start, end, 0, vtk.reference(0), temporaryHitPoint, [0.0, 0.0, 0.0], vtk.reference(0), vtk.reference(0))
                # if hit is found, and the top hitpoint is within ROI bounds
                if res != 0 and region_of_interest[0] <= temporaryHitPoint[castIndex] < region_of_interest[1]:
                    temporaryHitPoint[castIndex] += 0.3 * negated  # raised to improve visibility
                    hitPointMatrix[i][j] = HitPoint(points.InsertNextPoint(temporaryHitPoint), temporaryHitPoint[:])

        # form quads/cells
        update_status(text="Forming top layer polygons", progress=64)
        cells = vtk.vtkCellArray()
        for i in range(len(hitPointMatrix)-1):  # -1 as the end row/col will be taken into account
            for j in range(len(hitPointMatrix[i])-1):
                hitPoints = [hitPointMatrix[i][j], hitPointMatrix[i+1][j], hitPointMatrix[i+1][j+1], hitPointMatrix[i][j+1]]
                # check if full quad
                if None in hitPoints: continue
                # check if area is not extremely large
                d1 = numpy.linalg.norm(numpy.array(hitPoints[0].point)-numpy.array(hitPoints[1].point))
                d2 = numpy.linalg.norm(numpy.array(hitPoints[0].point)-numpy.array(hitPoints[2].point))
                d3 = numpy.linalg.norm(numpy.array(hitPoints[0].point)-numpy.array(hitPoints[3].point))
                m = precision*6
                if d1 > m or d2 > m or d3 > m: continue
                # calculate normals
                rawNormal = numpy.linalg.solve(numpy.array([hitPoints[0].point, hitPoints[1].point, hitPoints[2].point]), [1, 1, 1])
                hitPointMatrix[i][j].normal = rawNormal / numpy.sqrt(numpy.sum(rawNormal**2))
                # # check if quad is acceptable by normal vs. cast vector
                # v1, v2 = numpy.array(hitPointMatrix[i][j].normal), numpy.array(castVector)
                # degrees = numpy.degrees(numpy.math.atan2(numpy.cross(v1, v2).shape[0], numpy.dot(v1, v2)))
                cells.InsertNextCell(4, [p.pid for p in hitPoints])
        update_status(text="Finished ray-casting in " + str("%.1f" % (time.time() - startTime)) + "s, found " + str(cells.GetNumberOfCells()) + " cells...", progress=80)

        # build poly data
        topLayerPolyData = vtk.vtkPolyData()
        topLayerPolyData.SetPoints(points)
        topLayerPolyData.SetPolys(cells)
        topLayerPolyData.Modified()
        return topLayerPolyData, [p for d0 in hitPointMatrix for p in d0 if p is not None]

    @staticmethod
    def build_model(poly_data, update_status):
        update_status(text="Rendering top layer...", progress=20)
        modelNode = slicer.mrmlScene.AddNewNodeByClass('vtkMRMLModelNode')
        modelNode.SetAndObservePolyData(poly_data)
        modelNode.CreateDefaultDisplayNodes()
        modelDisplayNode = modelNode.GetModelDisplayNode()
        modelDisplayNode.SetFrontfaceCulling(0)
        modelDisplayNode.SetBackfaceCulling(0)
        update_status(text="Top layer rendered...", progress=40)
        return modelNode

    @staticmethod
    def ray_cast_color_thickness(poly_data, hit_point_list, cast_axis, dimensions, mm_of_air_past_bone, update_status, gradient_scale_factor=10.0):
        # ray direction cast axis index
        castIndex = BoneThicknessMappingLogic.determine_cast_axis_index(cast_axis)

        update_status(text="Building static cell locator...", progress=81)
        cellLocator = vtk.vtkStaticCellLocator()
        cellLocator.SetDataSet(poly_data)
        cellLocator.BuildLocator()

        total = len(hit_point_list)
        update_status(text="Calculating thickness (may take long, " + str(total) + " rays)...", progress=82)
        startTime = time.time()

        def init_array(name):
            a = vtk.vtkFloatArray()
            a.SetName(name)
            return a

        def interpret_distance(points):
            firstIn, lastOut = points[0], points[1]
            if len(points) > 2:
                for inOutPair in zip(*[iter(points[2:])] * 2):
                    # TODO incorporate cast bounds
                    # if newIn[1][cast_axis] < minBound and newOut[1][cast_axis] < minBound
                    if calculate_distance(lastOut[1], inOutPair[0][1]) < mm_of_air_past_bone*gradient_scale_factor: lastOut = inOutPair[1]
                return calculate_distance(firstIn[1], lastOut[1])
            else:
                return calculate_distance(firstIn[1], points[-1][1])

        def calculate_distance(point1, point2):
            d = numpy.linalg.norm(numpy.array((point1[0], point1[1], point1[2])) - numpy.array((point2[0], point2[1], point2[2])))
            d = d*gradient_scale_factor
            return d

        skullThicknessArray, airCellDistanceArray = init_array(BoneThicknessMappingType.THICKNESS), init_array(BoneThicknessMappingType.AIR_CELL)
        tol, pCoords, subId = 0.000, [0, 0, 0], vtk.reference(0)
        pointsOfIntersection, cellsOfIntersection = vtk.vtkPoints(), vtk.vtkIdList()
        for i, hitPoint in enumerate(hit_point_list):
            stretchFactor = dimensions[castIndex]
            start = [hitPoint.point[0] + hitPoint.normal[0]*stretchFactor, hitPoint.point[1] + hitPoint.normal[1]*stretchFactor, hitPoint.point[2] + hitPoint.normal[2]*stretchFactor]
            end = [hitPoint.point[0] - hitPoint.normal[0]*stretchFactor, hitPoint.point[1] - hitPoint.normal[1]*stretchFactor, hitPoint.point[2] - hitPoint.normal[2]*stretchFactor]
            cellLocator.FindCellsAlongLine(start, end, tol, cellsOfIntersection)
            distances = []
            for cellIndex in range(cellsOfIntersection.GetNumberOfIds()):
                t = vtk.reference(0.0)
                p = [0.0, 0.0, 0.0]
                if poly_data.GetCell(cellsOfIntersection.GetId(cellIndex)).IntersectWithLine(start, end, tol, t, p, pCoords, subId) and 0.0 <= t <= 1.0:
                    distances.append([t, p])
            if len(distances) >= 2:
                distances = sorted(distances, key=lambda kv: kv[0])
                skullThicknessArray.InsertTuple1(hitPoint.pid, interpret_distance(distances))
                airCellDistanceArray.InsertTuple1(hitPoint.pid, calculate_distance(distances[0][1], distances[1][1]))
            else:
                skullThicknessArray.InsertTuple1(hitPoint.pid, 0)
                airCellDistanceArray.InsertTuple1(hitPoint.pid, 0)
            # update rays casted status
            if i % 200 == 0: update_status(text=f"Calculating thickness (~{i} of {total} rays)", progress=82 + int(round((i*1.0/total*1.0)*18.0)))
        update_status(text="Finished thickness calculation in " + str("%.1f" % (time.time() - startTime)) + "s...", progress=100)
        return skullThicknessArray, airCellDistanceArray

    @staticmethod
    def build_color_table_node(name, table_max):
        table = slicer.vtkMRMLColorTableNode()
        table.SetName(name)
        table.SetHideFromEditors(0)
        table.SetTypeToFile()
        table.NamesInitialisedOff()
        table.SetNumberOfColors(table_max)
        table.GetLookupTable().SetTableRange(0, table_max)
        table.NamesInitialisedOn()
        slicer.mrmlScene.AddNode(table)
        return table

    @staticmethod
    def build_color_table_nodes(minmax_thickness, minmax_air_cell, gradient_scale_factor=10.0):
        print(minmax_air_cell)

        def calculate_and_set_colour(table, index, hue=0.0, sat=1.0, val=1.0):
            rgb = colorsys.hsv_to_rgb(hue, sat, val)
            table.SetColor(index, str(index/gradient_scale_factor) + ' mm', rgb[0], rgb[1], rgb[2], 1.0)

        def p(lhs, rhs, base, right_mod=1.0):
            return float(lhs-base)/float((rhs-base)*right_mod)

        # thickness table
        ix = [int(i) for i in [minmax_thickness[0]*gradient_scale_factor, (minmax_thickness[1])*gradient_scale_factor + 1]]
        thicknessTableNode = BoneThicknessMappingLogic.build_color_table_node('ThicknessColorMap', ix[-1])
        for i in range(ix[0], ix[-1]): calculate_and_set_colour(thicknessTableNode, i, hue=p(i, ix[-1], ix[0]) * 0.278, sat=0.9, val=0.9)

        # air cell table
        iy = [int(i) for i in [minmax_air_cell[0]*gradient_scale_factor, (minmax_air_cell[1])*gradient_scale_factor + 1]]
        airCellTableNode = BoneThicknessMappingLogic.build_color_table_node('AirCellColorMap', iy[-1])
        for i in range(iy[0], iy[-1]): calculate_and_set_colour(airCellTableNode, i, hue=0.696 - p(i, iy[-1], iy[0]) * 0.571, sat=0.9, val=0.9)

        return thicknessTableNode, airCellTableNode

    @staticmethod
    def set_scalar_colour_bar_state(state, color_node_id=None):
        colorWidget = slicer.modules.colors.widgetRepresentation()
        ctkBar = slicer.util.findChildren(colorWidget, name='VTKScalarBar')[0]
        ctkBar.setDisplay(state)
        if state == 0 or color_node_id is None: return
        slicer.util.findChildren(colorWidget, 'ColorTableComboBox')[0].setCurrentNodeID(color_node_id)
        slicer.util.findChildren(colorWidget, 'UseColorNameAsLabelCheckBox')[0].setChecked(True)


from flask import Flask, request, render_template, jsonify
from pytube import YouTube
from pathlib import Path
import os

app = Flask(__name__)

@app.route("/")
def index():
    return render_template('index.html')

@app.route("/download", methods=["GET", "POST"])
def downloadVideo():
    message = ''
    errorType = 0

    if request.method == 'POST' and 'video_url' in request.form:
        youtubeUrl = request.form["video_url"]
        if youtubeUrl:
            url = YouTube(youtubeUrl)

            # 获取用户选择的清晰度和格式
            selected_quality = request.form['selected_stream']
            selected_stream = None

            # 根据用户选择的清晰度和格式筛选视频流
            for stream in url.streams:
                value = f"{stream}"
                if value == selected_quality:
                    selected_stream = stream
                    break

            if selected_stream:
                downloadFolder = str(os.path.join(Path.home(), "Downloads"))
                if selected_stream.resolution:
                    selected_stream.download(output_path=downloadFolder, filename_prefix=f"{selected_stream.resolution}_")
                else:
                    selected_stream.download(output_path=downloadFolder, filename_prefix=f"{selected_stream.abr}_")
                message = f'Video Downloaded Successfully in {downloadFolder}!'
                errorType = 0
            else:
                message = 'Selected stream is not available for download'
                errorType = 1
        else:
            message = 'Please enter YouTube Video URL.'
            errorType = 2
    return render_template('index.html', message=message, errorType=errorType)

@app.route("/get_quality", methods=["GET"])
def get_quality():
    video_url = request.args.get("url")
    if not video_url:
        return jsonify({"options": {"video": [], "audio": []}})

    url = YouTube(video_url)

    # 获取所有可用的视频流和音频流，将视频清晰度和格式对应起来
    available_streams = url.streams

    # 初始化 video 和 audio 的选项列表
    video_options = []
    audio_options = []

    # 将流按类型分类
    for stream in available_streams:
        if stream.type == "video":
            video_options.append({
            "value": f"{stream}",
            "res": f"{stream.resolution}",
            "subtype": f"{stream.subtype}",
            "label": f"{stream.resolution} - {stream.subtype}"
        })
        elif stream.type == "audio":
            audio_options.append({
            "value": f"{stream}",
            "abr": f"{stream.abr}",
            "subtype": f"{stream.subtype}",
            "label": f"{stream.abr} - {stream.subtype}"
        })

    # label去重
    def remove_duplicates(options, type):
        unique_options = []
        seen = set()
        for option in options:
            if option['label'] not in seen:
                unique_options.append(option)
                seen.add(option['label'])
        return unique_options

    video_options = sorted(remove_duplicates(video_options, type="video"), key=lambda x: (int(x["res"].replace('p', '')), x["subtype"]))

    audio_options = sorted(remove_duplicates(audio_options, type="audio"), key=lambda x: (int(x["abr"].replace('kbps', '')), x["subtype"]))

    return jsonify({"video": video_options, "audio": audio_options})

if __name__ == "__main__":
    app.run(host="0.0.0.0", port=5000, debug=True)