import sys, os
sys.path.insert(1, os.path.join(sys.path[0], '..'))
import importlib
import matplotlib
importlib.reload(matplotlib)
matplotlib.use('Qt4Agg')
import vispy
import numpy as np
import random
from PyQt4.QtGui import *
from PyQt4 import QtCore, QtGui

from vispy import app, gloo
from vispy.util.transforms import perspective, translate, rotate
from vispy.geometry import meshdata as md
from vispy.geometry import generation as gen

from Deadweight.ShapeFileHandler import ShapeFileHandler
from Core.LasFileHandler import LasFileHandler

from Utils.pca_utils import *

import matplotlib.pyplot as plt
from matplotlib.backends.backend_qt4agg import FigureCanvasQTAgg as FigureCanvas
from matplotlib.backends.backend_qt4agg import NavigationToolbar2QT

import Gui

from Core import DataHandler

class PlotObject():

    def __init__(self, parent=None):

        # a figure instance to plot on
        self.figure = plt.figure()

        # this is the Canvas Widget that displays the `figure`
        # it takes the `figure` instance as a parameter to __init__
        self.canvas = FigureCanvas(self.figure)

        # this is the Navigation widget
        # it takes the Canvas widget and a parent
        # self.toolbar = NavigationToolbar(self.canvas, self)
        self.toolbar = NavigationToolbar2QT(self.canvas, None)
        # Just some button connected to `plot` method
        self.button = QtGui.QPushButton('Plot')
        self.button.clicked.connect(self.plot)

        # set the layout
        self.layout = QtGui.QVBoxLayout()
        self.layout.addWidget(self.canvas)
        self.layout.addWidget(self.toolbar)
        self.layout.addWidget(self.button)
        #self.setLayout(layout)
        self.plot()

    def plot(self):
        ''' plot some random stuff '''
        # random data
        data = [random.random()*10 for i in range(10)]

        # create an axis
        ax = self.figure.add_subplot(111)

        # discards the old graph
        ax.hold(False)

        # plot data
        ax.plot(data, '*-')

        # refresh canvas
        self.canvas.draw()

class PipelineWidget(QWidget):

    def __init__(self, parent, data):

        QWidget.__init__(self)

        self.__parent__ = parent
        self.__data_handler = data
        self.current_data = {}
        self.setup_gui()
        self.resize(250, 150)

    def set_slice_from_property(self, property_id):
        self.current_data['bfps'], self.current_data['points'], self.current_data['solids'], self.current_data['property_area'], self.current_data['mbb'], self.current_data['top_dogs'] = self.__data_handler.get_slice_from_property(property_id)
        #self.__offset = np.mean(self.__property_area.points,0)

        self.current_data['top_dog_points'] = [self.current_data['points'][i] for i in range(len(self.current_data['bfps'])) if self.current_data['bfps'][i].id in self.current_data['top_dogs'] ]

    def setup_gui(self):
        #DATA GROUP
        self.data_group = QGroupBox("Data")
        self.data_layout = QVBoxLayout()
        self.data_group.setMaximumWidth(250)
        self.data_group.setLayout(self.data_layout)

        #self.bfp_checkbox = QCheckBox("Bfp blocks")
        #self.bfp_checkbox.setChecked(False)
        #self.bfp_checkbox.stateChanged.connect(lambda: self.bfps(self.bfp_checkbox.isChecked() ))

        #RENDER GROUP
        self.rendering_group = QGroupBox("Rendering")
        self.rendering_layout = QVBoxLayout()
        self.rendering_group.setMaximumWidth(250)
        self.rendering_group.setLayout(self.rendering_layout)

        # render polygons
        self.render_polygons_checkbox = QCheckBox("Render polygons?")
        self.render_polygons_checkbox.setChecked(True)
        self.render_polygons_checkbox.stateChanged.connect(lambda: self.setDrawPolygons(self.render_polygons_checkbox.isChecked()))
        self.rendering_layout.addWidget(self.render_polygons_checkbox)
        # transparent checkbox
        self.transparent_checkbox = QCheckBox("Transparent polygons")
        self.transparent_checkbox.setChecked(False)
        self.transparent_checkbox.stateChanged.connect(lambda: self.transp(self.transparent_checkbox.isChecked()))
        self.rendering_layout.addWidget(self.transparent_checkbox)
        # top dogs checkbox
        self.topdogs_checkbox = QCheckBox("Top dogs only")
        self.topdogs_checkbox.setChecked(False)
        self.topdogs_checkbox.stateChanged.connect(lambda: self.dogs(self.topdogs_checkbox.isChecked()))
        self.rendering_layout.addWidget(self.topdogs_checkbox)

        # render points
        self.render_points_checkbox = QCheckBox("Render points?")
        self.render_points_checkbox.setChecked(True)
        self.render_points_checkbox.stateChanged.connect(lambda: self.setDrawPoints(self.render_points_checkbox.isChecked()))
        self.rendering_layout.addWidget(self.render_points_checkbox)
        # reset points
        self.reset_points_button = QPushButton("reset points")
        self.reset_points_button.clicked.connect(lambda: self.reset_points())
        self.rendering_layout.addWidget(self.reset_points_button)
        #Colors
        self.colors_checkbox = QCheckBox("Debug colors")
        self.colors_checkbox.setChecked(False)
        self.colors_checkbox.stateChanged.connect(lambda: self.colors(self.colors_checkbox.isChecked()))
        self.rendering_layout.addWidget(self.colors_checkbox)

        #CLUSTERING GROUP
        self.cluster_group = QGroupBox("Clustering")
        self.cluster_layout = Gui.ClusterLayout(self)
        self.cluster_group.setMaximumWidth(300)
        self.cluster_group.setLayout(self.cluster_layout)

        #PLANE FITTING GROUP
        self.plane_group = QGroupBox("Plane Fitting")
        self.plane_layout = Gui.PlaneFittingLayout(self)
        self.plane_group.setMaximumWidth(200)
        self.plane_group.setLayout(self.plane_layout)

        #PLANE FITTING GROUP
        self.features_group = QGroupBox("Roof features creation")
        self.features_layout = Gui.FeatureExtractionLayout(self)
        self.features_group.setMaximumWidth(300)
        self.features_group.setLayout(self.features_layout)

        self.parent_layout = QHBoxLayout()
        self.pipeline_layout = QHBoxLayout()

        #self.pipeline_layout.addWidget(self.data_group)
        self.pipeline_layout.addWidget(self.rendering_group)
        self.pipeline_layout.addWidget(self.cluster_group)
        self.pipeline_layout.addWidget(self.plane_group)
        self.pipeline_layout.addWidget(self.features_group)

        self.parent_layout.setAlignment(QtCore.Qt.AlignLeft)
        self.parent_layout.addLayout(self.pipeline_layout)

        self.setLayout(self.parent_layout)

    def setDrawPoints(self, bool):
        self.__parent__.canvas_3d.render_params['draw_points'] = bool
        self.__parent__.canvas_3d.updateCanvas()

    def setDrawPolygons(self, bool):
        self.__parent__.canvas_3d.render_params['draw_polygons'] = bool
        self.__parent__.canvas_3d.updateCanvas()

    def setRegions(self, regions, topregions):
        self.current_data['regions'] = regions
        self.current_data['top_dog_regions'] = topregions
        self.__parent__.canvas_3d.regions()

    def reset_points(self):
        self.__parent__.canvas_3d.points()

    def setPlanes(self, planes, topplanes):
        self.current_data['planes'] = planes
        self.current_data['top_dog_planes'] = topplanes
        self.__parent__.canvas_3d.planes()

    def colors(self, value):
        self.__parent__.canvas_3d.render_params['debug_colors'] = value

    def setFeatures(self, features, topfeatures):
        self.current_data['features'] = features
        self.current_data['top_dog_features'] = topfeatures

        # import mpl_toolkits.mplot3d as a3
        # import matplotlib.colors as colors
        # import pylab as pl
        # import scipy as sp
        # import triangle
        # ax = a3.Axes3D(pl.figure())
        # colors = [[1,0,0],[0,1,0], [0,0,1], [0.65, 0.65, 0], [0, 0.65, 0.65], [0.65, 0, 0.65], [0.43, 0.43, 0.43]]
        #
        #
        # f = features
        # if f == None:
        #     f = topfeatures
        #
        # cid = 0
        # for feature_group in f:
        #     for feature in feature_group:
        #         roofPolys = feature["roof"].get_triangles_as_polygons()
        #         wallPolys = feature["walls"].get_triangles_as_polygons()
        #         #vertexdata = feature.get_gl_vertex_data()
        #
        #         tri = a3.art3d.Poly3DCollection([poly.exterior.coords for poly in wallPolys])
        #         tri.set_facecolor([0.8,0.8, 0.8])
        #         ax.add_collection3d(tri)
        #         tri = a3.art3d.Poly3DCollection([poly.exterior.coords for poly in roofPolys])
        #         tri.set_facecolor(colors[cid%7])
        #
        #         ax.add_collection3d(tri)
        #     cid += 1
        # ax.view_init(elev=39, azim=-55)
        # ax.set_xlim3d(-30, 30)
        # ax.set_ylim3d(-30, 30)
        # ax.set_zlim3d(0, 60)
        # pl.show()

        self.__parent__.canvas_3d.features()

    def transp(self, bool):
        self.__parent__.canvas_3d.render_params['transparent'] = bool

    def dogs(self, bool):
        self.__parent__.canvas_3d.render_params['top_dogs_only'] = bool
        self.__parent__.canvas_3d.updateCanvas()

    def bfps(self, bool):
        self.__parent__.canvas_3d.draw_model_transparent = bool

    def render_points(self, attr, value):
        #self.__parent__.canvas_3d.setup_bfps()
        pass

class MainWindow(QMainWindow):
    """
        this is the mainWondow for the program gui
    """

    def __init__(self, data_handler):

        QMainWindow.__init__(self)

        self.__data_handler = data_handler

        self.resize(1400, 840)
        self.setWindowTitle('gui main window')

        self.pipelineWidget = PipelineWidget(self, data_handler)

        self.canvas_map = Gui.Map2D(data_handler, self)
        self.canvas_map.create_native()
        self.canvas_map.native.setParent(self)

        self.canvas_3d = Gui.Canvas3D(data_handler, self, self.pipelineWidget.current_data)
        self.canvas_3d.create_native()
        self.canvas_3d.native.setParent(self)

        self.listWidget = Gui.BuildingListWidget(self, data_handler)

        self.setupLayout()
        #self.__focus_building_id = 0

        self.set_slice(1592) #1593
        self.show()

    #Is called from the property list
    def set_slice(self, id):
        print("selected property id: ", id)
        self.pipelineWidget.set_slice_from_property(id)
        self.canvas_3d.setCanvasData()

    @property
    def focus_building_id(self):
        return self.__focus_building_id

    @focus_building_id.setter
    def focus_building_id(self, id_to_focus):
        self.__focus_building_id = id_to_focus
        self.canvas_map.set_selected(self.__focus_building_id)
        #self.canvas_3d.set_selected(self.__focus_building_id)

    def setupLayout(self):

        self.verticalSplitter = QSplitter(QtCore.Qt.Vertical)

        self.horizontalSplitter = QSplitter(QtCore.Qt.Horizontal)

        self.mapListSplitter = QSplitter(QtCore.Qt.Vertical)

        self.mapListSplitter.addWidget(self.listWidget)
        self.mapListSplitter.addWidget(self.canvas_map.native)

        self.horizontalSplitter.addWidget(self.mapListSplitter)

        self.horizontalSplitter.addWidget(self.canvas_3d.native)

        self.verticalSplitter.addWidget(self.horizontalSplitter)
        self.verticalSplitter.addWidget(self.pipelineWidget)
        self.setCentralWidget(self.verticalSplitter)

if __name__ == "__main__":
    appQt = QApplication(sys.argv)

    #value = Vestibule(None)

    #filename = ("by_get_shp/by_get")

    #shapes = ShapeFileHandler(filename)

    #las_dir_path = os.path.dirname(os.path.realpath(__file__))
    #las_dir_path = las_dir_path + '/datafromndr'
    #las_file_path = las_dir_path + "/norrkoping.las"
    #las = LasFileHandler(las_file_path,8)

    #las.read_points()
    #las.read_points()

    #mw = MainWindow(shapes, las) #gui main

    buildings_path = "by_get_shp/by_get"
    propertys_path = "ay_get_shp/ay_get"
    las_path = "datafromndr/norrkoping.las"
    data_handler = DataHandler(buildings_path,propertys_path,las_path)

    mw = MainWindow(data_handler)

    sys.exit(appQt.exec_())


import torch

def evaluate_gan(generator, test_loader, device):
    generator.eval()
    total_loss = 0
    criterion = torch.nn.MSELoss()

    with torch.no_grad():
        for images, voxels in test_loader:
            images = images.to(device)
            voxels = voxels.to(device)
            generated_voxels = generator(images)
            loss = criterion(generated_voxels, voxels)
            total_loss += loss.item()

    generator.train()
    return total_loss / len(test_loader)


import torch
from interactive_navigation.utils.train_mask_rcnn import get_model_instance_segmentation
from PIL import Image
import numpy as np
import cv2

if __name__ == "__main__":
    root_dir = "storage/heuristic_corner_detector_test_scene"
    OBJ = sorted(["ArmChair", "DogBed", "Box", "Chair", "Desk", "DiningTable", "SideTable", "Sofa",
                  "Stool", "Television", "Pillow", "Bread", "Apple", "AlarmClock", "Lettuce",
                  "GarbageCan", "Laptop", "Microwave", "Pot", "Tomato"])
    SIZE = (400, 400)
    COLOR = [(0, 0, 0),
             (255, 0, 0),
             (255, 128, 0),
             (255, 255, 0),
             (128, 255, 0),
             (0, 255, 0),
             (0, 255, 128),
             (0, 255, 255),
             (0, 128, 255),
             (0, 0, 255),
             (127, 0, 255),
             (255, 0, 255),
             (255, 0, 127),
             (128, 128, 128),
             (102, 102, 0),
             (102, 0, 0),
             (0, 102, 0),
             (0, 102, 102),
             (0, 0, 102),
             (102, 0, 102),
             (255, 255, 255)]

    gpu_id = 0
    loc = 'cuda:{}'.format(gpu_id)
    checkpoint = torch.load("storage/maskRcnn/model_9.pth", map_location=loc)
    model = get_model_instance_segmentation(21).cuda(gpu_id)
    model.load_state_dict(checkpoint["model"])
    model.eval()

    imgs = []
    for obj in OBJ:
        img = Image.open("{}/{}/obs.png".format(root_dir, obj))
        img = img.resize(SIZE)
        img = np.array(img) / 255.
        img = np.transpose(img, (2, 0, 1))
        img = torch.Tensor(img).to(gpu_id)
        imgs.append(img)
    output = model(imgs)
    for j, out in enumerate(output):
        labels = out["labels"].detach().cpu().numpy()
        scores = out["scores"].detach().cpu().numpy()
        masks = out["masks"].squeeze().detach().cpu().numpy()
        result = np.zeros((SIZE[0], SIZE[1], 3))
        for i, label in enumerate(labels):
            if label == (j + 1):
                if scores[i] > 0.5:
                    color = COLOR[label]
                    mask = masks[i]
                    idx = np.where(mask > 0.1)
                    result[idx] = color
        result = result.astype(np.uint8)
        img = Image.fromarray(result)
        img.save("{}/{}/maskrcnn_segmentation.png".format(root_dir, OBJ[j]))

        org_img = np.transpose(imgs[j].detach().cpu().numpy() * 255, (1, 2, 0)).astype(np.uint8)
        fin = cv2.addWeighted(result, 0.5, org_img, 0.85, 0)
        fin = Image.fromarray(fin)
        fin.save("{}/{}/maskrcnn_segmentation_blend.png".format(root_dir, OBJ[j]))

        org_img = Image.fromarray(org_img)
        org_img.save("{}/{}/maskrcnn_org_image.png".format(root_dir, OBJ[j]))

    xxx = 0


from sklearn.cluster import  MiniBatchKMeans, KMeans, DBSCAN
import numpy as np
import matplotlib.pyplot as plt
import scipy.io
import pandas as pd
import math
from pylab import *
from scipy.spatial.distance import cdist
from scipy import spatial
from mpl_toolkits.mplot3d import Axes3D
import time

from utils.logistics_functions import *
from utils.Parameters_class import Parameters_values

Parameters = Parameters_values()

class logistics_model():
	def __init__(self, path, n_regen, n_collection, logistics, scenario):
		self.path =path
		self.n_regen = n_regen
		self.n_collection = n_collection
		self.logistics = logistics
		self.scenario = scenario

	def return_dataframe_buildings(self):
		building_virtual_buildings_df = pd.read_csv(self.path)
		building_SF_points=building_virtual_buildings_df[['lat_lat','lon_lon']]
		return building_SF_points

	def clustering_regen(self):
		k_means_labels_regen, k_means_cluster_centers_regen = clustering(self.return_dataframe_buildings(), self.n_regen)
		return k_means_labels_regen, k_means_cluster_centers_regen

	def clustering_regen_grid(self):
		building_SF_points = self.return_dataframe_buildings()
		nx = int(sqrt(self.n_regen))
		ny = int(ceil(self.n_regen/nx))
		cluster_centers_regen = make_grid_points(self.n_regen, nx, ny)
		cluster_centers_regen_df=pd.DataFrame(cluster_centers_regen, columns=['lat', 'lon'])
		tree = spatial.KDTree(list(zip(cluster_centers_regen_df['lat'], cluster_centers_regen_df['lon'])))
		building_SF_points_array = list(zip(building_SF_points['lat_lat'],building_SF_points['lon_lon']))
		k_means_labels_regen = tree.query(building_SF_points_array)[1]

		return k_means_labels_regen, cluster_centers_regen

	def clustering_regen_random(self):
		building_SF_points = self.return_dataframe_buildings()
		cluster_centers_regen = make_random_points(self.n_regen)
		cluster_centers_regen_df=pd.DataFrame(cluster_centers_regen, columns=['lat', 'lon'])
		tree = spatial.KDTree(list(zip(cluster_centers_regen_df['lat'], cluster_centers_regen_df['lon'])))
		building_SF_points_array = list(zip(building_SF_points['lat_lat'],building_SF_points['lon_lon']))
		k_means_labels_regen = tree.query(building_SF_points_array)[1]

		return k_means_labels_regen, cluster_centers_regen

	def clustering_collection(self, k_means_cluster_centers_regen):
		k_means_labels_collection, k_means_cluster_centers_collection = clustering(k_means_cluster_centers_regen, self.n_collection)
		return k_means_labels_collection, k_means_cluster_centers_collection

	def logistics_distances(self):
		building_virtual_buildings_df = pd.read_csv(self.path)

		if self.logistics == 'optimal':
			k_means_labels_regen, k_means_cluster_centers_regen = self.clustering_regen()

		elif self.logistics == 'grid':
			k_means_labels_regen, k_means_cluster_centers_regen = self.clustering_regen_grid()

		elif self.logistics == 'random':
			k_means_labels_regen, k_means_cluster_centers_regen = self.clustering_regen_random()

		else:
			print ("Error type of logistics was not specified")

		time_now = time.time()

		print ("Start calculating distances...")
		distance_regeneration = find_distance_regeneration_scheduled(building_virtual_buildings_df, k_means_labels_regen, self.scenario)
		distance_regeneration['total_dist_m_y'] = distance_regeneration['total_dist_m']*365/Parameters.time_between_catridge_regeneration_2
		time_end = time.time() - time_now
		print ("calc distances took time %s" %time_end)

		distance_collection = find_distance_collection(k_means_cluster_centers_regen)*Parameters.collection_times_per_year
		return distance_regeneration, distance_collection


from ..common.boundingbox import BoundingBox
import cv2
import numpy as np

""" 各自圖片進行撿均值（突顯該 圖像 重點） """
def subtract_avg(img):

    trg_img = img
    for i in range(3):
        avg=np.average(trg_img[:,:,i])
        trg_img[:,:,i]-=avg
    return trg_img

""" 整個數據集進行撿均值（突顯該 類別 重點） """
def subtract_offset(img, offset=( 103.939, 116.779, 123.68 )):
    trg_img = img
    for i in range(3):
        trg_img[:,:,i]=img[:,:,i]-offset[i]
    return trg_img

""" caffe mode of image processing """
def preprocess(img, input_shape, letter_box=False):
    """Preprocess an image before TRT YOLO inferencing.
    # Args
        img: int8 numpy array of shape (img_h, img_w, 3)
        input_shape: a tuple of (H, W)
        letter_box: boolean, specifies whether to keep aspect ratio and
                    create a "letterboxed" image for inference
    # Returns
        preprocessed img: float32 numpy array of shape (3, H, W)
    """
    if letter_box:
        img_h, img_w, _ = img.shape
        new_h, new_w = input_shape[0], input_shape[1]
        offset_h, offset_w = 0, 0
        if (new_w / img_w) <= (new_h / img_h):
            new_h = int(img_h * new_w / img_w)
            offset_h = (input_shape[0] - new_h) // 2
        else:
            new_w = int(img_w * new_h / img_h)
            offset_w = (input_shape[1] - new_w) // 2
        resized = cv2.resize(img, (new_w, new_h))
        img = np.full((input_shape[0], input_shape[1], 3), 127, dtype=np.uint8)
        img[offset_h:(offset_h + new_h), offset_w:(offset_w + new_w), :] = resized
    else:
        img = cv2.resize(img, (input_shape[1], input_shape[0]))

    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    img = img.transpose((2, 0, 1)).astype(np.float32)
    img /= 255.0
    return img

def postprocess(output, img_w, img_h, input_shape, conf_th=0.8, nms_threshold=0.5, letter_box=False):
    """
    (-1,200), (-1,200,4), (-1,200), (-1,200)
    num_detections: A [batch_size] tensor containing the INT32 scalar indicating the number of valid detections per batch item. It can be less than keepTopK. Only the top num_detections[i] entries in nmsed_boxes[i], nmsed_scores[i] and nmsed_classes[i] are valid
    nmsed_boxes: A [batch_size, keepTopK, 4] float32 tensor containing the coordinates of non-max suppressed boxes
    nmsed_scores: A [batch_size, keepTopK] float32 tensor containing the scores for the boxes
    nmsed_classes: A [batch_size, keepTopK] float32 tensor containing the classes for the boxes
    """
    (detections, bboxes, scores, labels) = output

    results= []
    for idx, det in enumerate(detections):
        # print(scores[idx], '\t')
        x1, y1, x2, y2 = map(float, bboxes[idx]) # x1, y1, x2, y2
        x1, y1, x2, y2 = x1*img_w, y1*img_h, x2*img_w, y2*img_h
        if scores[idx]>=conf_th:
            results.append(BoundingBox(labels[idx], scores[idx], x1, x2, y1, y2, img_h, img_w))
    return results


import torch
import torch.optim as optim

import numpy as np
import time

import os
from datetime import datetime
from torch.utils.tensorboard import SummaryWriter

from Seq2Seq_model import Seq2Seq_VAE

from helper_utils import colorize
from training_functions import run_VAE

from helper_utils import random_trajectories

from LSL_utils import initial_region_points
from LSL_utils import find_good_region
from LSL_utils import set_param_grad_status

from LSL_utils import trajectory_diagnostics

import copy

import roboschool
import gym

def main(main_anneal='logistic',
         main_num_compression_iterations=30,
         main_num_logistic_compression_iterations=15,
         main_initial_num_epochs=3000,
         main_num_const_epochs=10,
         bootstrap="no", given_time="",
         final_model_file=None, old_final_model_file=None, argmax_only=False, small_steps=False,
         main_z_size=2, main_enc_hidden_size=12, main_dec_hidden_size=50, main_num_layers=1,
         main_bidir=True, main_use_layer_norm=True, main_use_ff_dropout=False, main_use_recurrent_dropout=False,
         main_ff_drop_prob=0.0, main_recurrent_drop_prob=0.0, main_dec_num_layers=3,
         main_Adam_learning_rate=5e-3, main_clip_gn=500.0, main_clip_gv=100.0, main_kl_tolerance=0.0,
         main_target_beta=1.0, main_logistic_scale=12.0, main_floor=0.5, main_num_grid_points=400,
         main_num_trajectories=240, main_batch_size=8, main_seq_len=20,
         main_env_name = "MyCartPoleNew-v0",
         main_torch_manual_seed=684, main_np_random_seed=684, main_env_seed=6328):

    torch.set_printoptions(precision=3, sci_mode=False)
    np.set_printoptions(precision=3, suppress=True)

    torch_manual_seed = main_torch_manual_seed
    np_random_seed = main_np_random_seed
    env_seed = main_env_seed

    torch.manual_seed(torch_manual_seed)
    np.random.seed(np_random_seed)

    env_name = main_env_name

    z_size = main_z_size
    enc_hidden_size = main_enc_hidden_size
    dec_hidden_size = main_dec_hidden_size
    num_layers = main_num_layers
    bidir = main_bidir
    use_layer_norm = main_use_layer_norm

    use_ff_dropout = main_use_ff_dropout
    use_recurrent_dropout = main_use_recurrent_dropout
    ff_drop_prob = main_ff_drop_prob
    recurrent_drop_prob = main_recurrent_drop_prob

    dec_num_layers = main_dec_num_layers

    seq_len = main_seq_len

    env = gym.make(env_name)
    env.seed(env_seed)
    obs_dim = env.observation_space.shape

    env_state_size = obs_dim[0]
    action_size = env.action_space.n

    Adam_learning_rate = main_Adam_learning_rate
    clip_gn = main_clip_gn
    clip_gv = main_clip_gv
    kl_tolerance = main_kl_tolerance
    target_beta = main_target_beta
    logistic_scale = main_logistic_scale

    anneal = main_anneal
    floor = main_floor

    num_grid_points = main_num_grid_points
    num_trajectories = main_num_trajectories
    batch_size = main_batch_size

    current_time = datetime.now().strftime('%b%d_%H-%M-%S')

    model_file = ('saved_models/model' + anneal + str(torch_manual_seed) + current_time + '.pt')

    num_compression_iterations = main_num_compression_iterations

    trajectory_actions, trajectory_actions_one_hot, trajectory_env_states = \
        random_trajectories(num_trajectories, seq_len, action_size, env_state_size, env=env)

    primitive_old_model_file = old_final_model_file
    primitive_model_file = final_model_file

    if bootstrap != "no":
        load_model = True
    else:
        load_model = False

    if load_model:

        print("loading model: " + primitive_model_file)
        model: Seq2Seq_VAE = torch.load(primitive_model_file)
        print(model)

        new_model = copy.deepcopy(model)

        if primitive_old_model_file is not None:
            print("loading model: " + primitive_old_model_file)
            model: Seq2Seq_VAE = torch.load(primitive_old_model_file)
            print(model)

        new_model.h2_bottom_seq_len = 7
        new_model.h2_top_seq_len = 10

        new_model.h3_bottom_seq_len = 7
        new_model.h3_top_seq_len = 10

    else:

        print("creating new model")

        num_hierarchies = 3

        model = Seq2Seq_VAE(action_size, env_state_size, z_size, enc_hidden_size, dec_hidden_size,
                            num_layers, dec_num_layers=dec_num_layers, seq_len=seq_len,
                            bidirectional_encoder=bidir,
                            forget_bias=1.0,
                            use_ff_dropout=use_ff_dropout, ff_drop_prob=ff_drop_prob, same_ff_mask=True,
                            use_recurrent_dropout=use_recurrent_dropout, recurrent_drop_prob=recurrent_drop_prob,
                            same_recurrent_mask=True,
                            use_layer_norm=use_layer_norm, num_hierarchies=num_hierarchies)

        new_model = copy.deepcopy(model)

    device = "cpu"

    h1_encoder_params = []
    h1_decoder_params = []

    named_params = list(new_model.named_parameters())

    index = 0

    for name, param in named_params:

        print("i: " + str(index) + " ", name, param.shape)
        index += 1

        if "h1" in name:
            if "encode" in name:
                h1_encoder_params += [param]
            elif "decode" in name:
                h1_decoder_params += [param]
            else:
                h1_decoder_params += [param]

    h_param_groups = [{'params': h1_encoder_params, 'g_name': 'h1_encoder_params'},
                      {'params': h1_decoder_params, 'g_name': 'h1_decoder_params'},
                      ]

    new_optimizer = optim.AdamW(h_param_groups, lr=Adam_learning_rate, weight_decay=0.0)

    set_param_grad_status(new_optimizer, ['h1_encoder_params', 'h1_decoder_params'], [])

    if bootstrap == "no":
        num_epochs = main_initial_num_epochs
    else:
        num_epochs = anneal_logistic_epochs(0, bootstrap=bootstrap)

    current_time = datetime.now().strftime('%b%d_%H-%M-%S')

    z_recon_stats_file = ('z_recon_stats/' + 'full_' + current_time + '.txt')

    if bootstrap != "no":

        z_seq_len = seq_len
        z_grid = initial_region_points(num_grid_points, z_size)
        z_grid = torch.from_numpy(z_grid).float().to(device)

        if small_steps:
            pos_to_sample = torch.cat([z_grid]*4)
        else:
            pos_to_sample = z_grid

        trajectory_actions, trajectory_actions_one_hot, _, trajectory_env_states = \
            new_model.argmax_or_repulsion_primitive_from_region_bootstrap(pos_to_sample, z_grid, z_seq_len,
                                                            argmax=True, env=env,
                                                            bootstrap_level=bootstrap, old_model=model)

    for compress_iteration in range(num_compression_iterations):

        print(colorize("compress_iteration: ", 'red'), compress_iteration)
        print("num_epochs: ", num_epochs)

        log_dir = os.path.join('LSL_new_runs', current_time + '_' +
                               anneal + '_num_epochs_' + str(num_epochs) +
                               "comp_iter_" + str(compress_iteration) + '_' + str(torch_manual_seed))

        print(log_dir)

        run_VAE_statsfile = ('LSL_run_VAE_stats/' + 'full_' +
                             anneal + '_num_epochs_' + str(num_epochs) +
                             str(torch_manual_seed) + "_comp_iter_" + str(compress_iteration) + current_time + '.txt')

        print(colorize("Writing stats to: " + run_VAE_statsfile, 'green'))
        f = open(run_VAE_statsfile, 'a')

        writer = SummaryWriter(log_dir)

        if small_steps:
            num_epochs = 3
            anneal = 'const'

        run_VAE(new_model, new_optimizer, trajectory_env_states, trajectory_actions, trajectory_actions_one_hot,
                num_epochs=num_epochs, clip_gn=clip_gn, clip_gv=clip_gv, batch_size=batch_size,
                kl_tolerance=kl_tolerance, anneal=anneal, logistic_scale=logistic_scale, target_beta=target_beta,
                floor=floor, log_interval=1, log_epoch_interval=1, f=f, writer=writer)

        iteration_model_file = ('saved_models_full/model_h1'
                                + anneal + 'comp_iter' + str(compress_iteration) +
                                str(torch_manual_seed) + current_time + '.pt')

        torch.save(new_model, iteration_model_file)

        if compress_iteration == num_compression_iterations - 1:

            final_model_file = ('saved_models_full/model_h1_final' +
                                anneal + 'comp_iter' + str(compress_iteration) +
                                str(torch_manual_seed) + given_time + '.pt')

            torch.save(new_model, final_model_file)

        if compress_iteration == 0 or compress_iteration == num_compression_iterations - 1:
            f.close()
            writer.close()

        print(colorize("-----Finished VAE iteration-----\n", 'red'))

        z_grid = initial_region_points(num_grid_points, z_size)
        z_grid = torch.from_numpy(z_grid).float().to(device)

        z_seq_len = seq_len
        primitive_trajectory_len = 5

        num_samples = z_grid.shape[0]
        new_model.eval()

        pos_to_sample = find_good_region(new_model, z_grid, z_seq_len, primitive_trajectory_len,
                                         num_samples=num_samples, env=env,
                                         bootstrap=bootstrap, old_model=model)

        pos_to_sample = pos_to_sample[:num_samples]

        if small_steps:
            pos_to_sample = torch.cat([pos_to_sample]*4)

        print(colorize("-----Finished finding sample region-----\n", 'red'))
        new_model.eval()

        if bootstrap == "no":

            trajectory_actions_argmax, trajectory_actions_one_hot_argmax, \
                trajectory_probs_argmax, trajectory_states_argmax = \
                new_model.argmax_or_repulsion_primitive_from_region(pos_to_sample, z_seq_len, primitive_trajectory_len,
                                                                    argmax=True, env=env)
        else:
            trajectory_actions_argmax, trajectory_actions_one_hot_argmax, \
                trajectory_probs_argmax, trajectory_states_argmax = \
                new_model.argmax_or_repulsion_primitive_from_region_bootstrap(pos_to_sample, z_grid, z_seq_len,
                                                                              argmax=True, env=env,
                                                                              bootstrap_level=bootstrap,
                                                                              old_model=model)

        print(colorize("-----Finished argmax sampling-----\n", 'red'))

        start_time = time.time()

        if bootstrap == "no" and not small_steps:
            high_trajectory_actions_repulsion, high_trajectory_actions_one_hot_repulsion, \
            high_trajectory_probs_repulsion, high_trajectory_states_repulsion = \
                new_model.argmax_or_repulsion_primitive_from_region(pos_to_sample, z_seq_len, primitive_trajectory_len,
                                                                    low_repulsion_idx=7, high_repulsion_idx=9,
                                                                    argmax=False, env=env)

            mid_trajectory_actions_repulsion, mid_trajectory_actions_one_hot_repulsion, \
            mid_trajectory_probs_repulsion, mid_trajectory_states_repulsion = \
                new_model.argmax_or_repulsion_primitive_from_region(pos_to_sample, z_seq_len, primitive_trajectory_len,
                                                                    low_repulsion_idx=3, high_repulsion_idx=6,
                                                                    argmax=False, env=env)

            low_trajectory_actions_repulsion, low_trajectory_actions_one_hot_repulsion, \
            low_trajectory_probs_repulsion, low_trajectory_states_repulsion = \
                new_model.argmax_or_repulsion_primitive_from_region(pos_to_sample, z_seq_len, primitive_trajectory_len,
                                                                    low_repulsion_idx=0, high_repulsion_idx=2,
                                                                    argmax=False, env=env)
        elif not small_steps:
            high_trajectory_actions_repulsion, high_trajectory_actions_one_hot_repulsion, \
            high_trajectory_probs_repulsion, high_trajectory_states_repulsion = \
                new_model.argmax_or_repulsion_primitive_from_region_bootstrap(
                                                            pos_to_sample, z_grid, z_seq_len,
                                                            low_repulsion_idx=7, high_repulsion_idx=9,
                                                            argmax=False, env=env,
                                                            bootstrap_level=bootstrap, old_model=model)

            mid_trajectory_actions_repulsion, mid_trajectory_actions_one_hot_repulsion, \
            mid_trajectory_probs_repulsion, mid_trajectory_states_repulsion = \
                new_model.argmax_or_repulsion_primitive_from_region_bootstrap(pos_to_sample, z_grid, z_seq_len,
                                                                    low_repulsion_idx=3, high_repulsion_idx=6,
                                                                    argmax=False, env=env,
                                                                    bootstrap_level=bootstrap, old_model=model)

            low_trajectory_actions_repulsion, low_trajectory_actions_one_hot_repulsion, \
            low_trajectory_probs_repulsion, low_trajectory_states_repulsion = \
                new_model.argmax_or_repulsion_primitive_from_region_bootstrap(pos_to_sample, z_grid, z_seq_len,
                                                                    low_repulsion_idx=0, high_repulsion_idx=2,
                                                                    argmax=False, env=env,
                                                                    bootstrap_level=bootstrap, old_model=model)

        print("repulsion time: " + str(time.time() - start_time))
        print(colorize("-----Finished repulsion sampling-----\n", 'red'))

        num_trajectories = trajectory_actions_argmax.shape[0]

        trajectory_diagnostics(model=new_model, z_space_folder='z_space_trajectories', env=env,
                               epoch=compress_iteration,
                               z_grid=z_grid[:num_trajectories],
                               z_seq_len=z_seq_len, primitive_trajectory_len=5,
                               plot=True, show_plot=False,
                               z_recon_stats_file=z_recon_stats_file)

        if compress_iteration < main_num_logistic_compression_iterations and (argmax_only is False):

            num_epochs = anneal_logistic_epochs(compress_iteration, bootstrap="no")

            trajectory_actions = torch.cat([trajectory_actions_argmax,
                                            high_trajectory_actions_repulsion,
                                            mid_trajectory_actions_repulsion,
                                            low_trajectory_actions_repulsion])

            trajectory_actions_one_hot = torch.cat([trajectory_actions_one_hot_argmax,
                                                    high_trajectory_actions_one_hot_repulsion,
                                                    mid_trajectory_actions_one_hot_repulsion,
                                                    low_trajectory_actions_one_hot_repulsion])

            trajectory_env_states = torch.cat([trajectory_states_argmax,
                                               high_trajectory_states_repulsion,
                                               mid_trajectory_states_repulsion,
                                               low_trajectory_states_repulsion])
        else:

            num_epochs = main_num_const_epochs
            anneal = 'const'

            trajectory_actions = trajectory_actions_argmax
            trajectory_actions_one_hot = trajectory_actions_one_hot_argmax
            trajectory_env_states = trajectory_states_argmax

        if not small_steps:

            random_indices = np.arange(num_trajectories)
            np.random.shuffle(random_indices)

    print(colorize("-----Finished final VAE iteration-----\n", 'red'))
    new_model.eval()

    torch.save(new_model, model_file)
    new_model = torch.load(model_file)
    new_model.eval()

    return final_model_file

def anneal_logistic_epochs(compress_iteration, bootstrap="no"):

    if bootstrap == "no":
        if compress_iteration < 5:
            return 5
        if compress_iteration < 10:
            return 8
        if compress_iteration < 15:
            return 8

        return 8
    else:
        return 5


import svgwrite

def inches(coords):
  a = "%1.4fin" % coords[0]
  b = "%1.4fin" % coords[1]
  return (a,b)

def draw_line(dwg, start, end, stroke=svgwrite.rgb(10,10,16,'%')):
    start = inches(start)
    end = inches(end)
    stroke = stroke
    print("Start at %s, End at %s" % (start, end))
    dwg.add(dwg.line(start, end, stroke=stroke))

# all widths are in inches
canvas_width = "20in"
canvas_height = "12in"

# Draftboard dividers
material = "draftboard"
slot_width = 0.124

# Acrylic dividers
#material = "acrylic"
#slot_width = 0.108

dd_length = 16.7 # same as width
dd_height = 5.38
slot_height = dd_height / 2
end_spacer = .5  # how wide is the last tab on each side?
num_of_slots = 7
slot_spacing = (-(slot_width * num_of_slots)-(2*end_spacer)+dd_length)/(num_of_slots-1)

filename = ("%s drawer_divider %sx%s - %d slots - %s slot width - %1.4f spacing - %s end tabs.svg"
            % (material, dd_length, dd_height, num_of_slots, slot_width, slot_spacing, end_spacer))
# offset to start drawing
x = 1
y = 1

dwg = svgwrite.Drawing(filename, profile='tiny', height=canvas_height, width=canvas_width)

lines = [
    ((x,y), (x,y+dd_height)), # Left side vertical line
    ((x+dd_length,y), (x+dd_length,y+dd_height)), # right side, vertical line
    ((x,y+dd_height), (x+dd_length, y+dd_height)), # bottom edge
    ((x,y), (x+end_spacer, y)), # left side, first tab
    ((x+dd_length, y), (x+dd_length-end_spacer, y)), # right side, last tab
]

for line in lines:
    draw_line(dwg, line[0], line[1])

# offset the width of the first tab
xx = x+end_spacer

print("draw first slot")
draw_line(dwg, (xx, y), (xx, y+slot_height))  # down
draw_line(dwg, (xx, y+slot_height), (xx+slot_width, y+slot_height))  # over
draw_line(dwg, (xx+slot_width, y+slot_height), (xx+slot_width, y))  # up
xx = xx + slot_width

for i in range(1,num_of_slots):
    print("draw %s slot" % (i+1))
    draw_line(dwg, (xx, y), (xx+slot_spacing, y)) # tab
    xx = xx+slot_spacing
    draw_line(dwg, (xx, y), (xx, y+slot_height)) # down
    draw_line(dwg, (xx,y+slot_height), (xx+slot_width, y+slot_height)) # over
    draw_line(dwg, (xx+slot_width, y+slot_height), (xx+slot_width, y)) # up
    xx = xx+slot_width

dwg.save()


'''sexo = ''
while sexo != 'M' and sexo != 'F':
    sexo = str(input('Informe seu sexo: [M/F] ')).strip().upper()
    if sexo != 'M' and sexo != 'F':
        print('Digite a letra correta!!!')
if sexo == 'M':
    sexo = 'Masculino'
else:
    sexo = 'Feminino'
print('O sexo escolhido foi {}'.format(sexo))
'''

sexo = str(input('Informe seu sexo: [M/F] ')).strip().upper()[0]
while sexo not in 'MmFf':
    sexo = str(input('Dados inválidos.Por favor, informe seu sexo novamente: ')).strip().upper()[0]
if sexo == 'M' or sexo == 'm':
    sexo = 'Masculino'
else:
    sexo = 'Feminino'
print('Sexo {} registrado com sucesso'.format(sexo))


import statsmodels.api as sm
import matplotlib.pyplot as plt
import numpy as np
import itertools
from scipy.stats import skewnorm

means = [100, 105]
std_devs = [5, 6]

num_to_sim = 500

data = [
    np.random.randn(num_to_sim) * s + mu for mu, s in zip(means, std_devs)
] + [skewnorm.rvs(10, size=num_to_sim)*std_devs[0] + means[0]]

plt.hist(data[0])
plt.xlabel('Value')
plt.ylabel('Count')
plt.savefig('imgs/basic_histogram.png')

plt.figure()
plt.hist(data[0])
plt.hist(data[1], alpha=0.5)
# plt.hist(data[2], alpha=0.5)

plt.xlabel('Value')
plt.ylabel('Count')
plt.savefig('imgs/two_distributions.png')

pp = sm.ProbPlot(data[0])
pp2 = sm.ProbPlot(data[1])
pp3 = sm.ProbPlot(data[2])
fig = pp.probplot()
pp2.probplot(ax=fig.gca())
pp3.probplot(ax=fig.gca())
plt.grid(True)

colors = itertools.cycle(plt.rcParams['axes.prop_cycle'].by_key()['color'])

for l in fig.gca().get_lines():
    c = next(colors)
    l.set_markerfacecolor(c)
    l.set_color(c)

plt.tight_layout()
plt.savefig('imgs/probplot.png')


def solution(m, musicinfos):
    answer = []
    m = m.replace('C#', 'c').replace('D#', 'd').replace('F#', 'f').replace('A#', 'a').replace('G#', 'g')
    for info in musicinfos:
        start, end, name, music = info.split(',')
        h, mm = start.split(':')
        sTime = int(h) * 60 + int(mm)
        h, mm = end.split(':')
        eTime = int(h) * 60 + int(mm)
        playTime = eTime - sTime
        tmp = ''
        i = 0
        music = music.replace('C#', 'c').replace('D#', 'd').replace('F#', 'f').replace('A#', 'a').replace('G#', 'g')
        for _ in range(playTime):
            tmp += music[i]
            i = (i + 1) % len(music)

        if m in tmp:
            answer.append([name, playTime])

    answer.sort(key= lambda x:x[1], reverse=True)

    return answer[0][0] if answer else "(None)"

프로그래머스/lv2/17683. ［3차］ 방금그곡/［3차］ 방금그곡.py

［3차］ 방금그곡.py


def comb(li, n):
    result = []
    if n > len(li):
        return result

    if n == 1:
        for i in li:
            result.append([i])
    else:
        for i in range(len(li) - n + 1):
            for j in comb(li[i+1:], n-1):
                result.append([li[i]] + j)

    return result

N, M = map(int, input().split())

num_li = []

for i in range(1, N+1):
    num_li.append(i)
# 조합 출력
ans = comb(num_li, M)
for v in ans:
    print(*v)


#!/usr/bin/env python
# coding: utf-8

'''
File: cron.py
Project: sjtuannouncements
File Created: 2019/06/03 15:30
Author: Roger Cui (itabas016@gmail.com)
Last Modified: 2019/06/03 15:31
Copyright 2019
'''

from celery.schedules import crontab
from celery.task import periodic_task

from sjtuannouncements import create_app
from sjtuannouncements.task import *
from sjtuannouncements.task import task

app = create_app('env', False)
celery = create_celery(app)

@periodic_task(run_every=(crontab(minute=0, hours='*/6,8-20')), name="scheduled_refresh_announcements_task", ignore_result=True)
def scheduled_refresh_announcements_task():
    with app.app_context():
        task.refresh_announcements()

@celery.task(bind=True)
def refresh_announcements_task(self):
    def listener(current, total, result):
        self.update_state(state='PROGRESS', meta={'progress': current / total})

    with app.app_context():
        return {'progress': 1, 'result': task.refresh_announcements(listener)}


import datetime
import sys
from gettext import gettext as _

from gi.repository import Gtk, Gio, GLib, Notify, Adw, GdkPixbuf, Gdk, GObject

from frog.config import RESOURCE_PREFIX, APP_ID
from frog.language_manager import language_manager
from frog.services.clipboard_service import clipboard_service
from frog.services.screenshot_service import ScreenshotService
from frog.settings import Settings
from frog.window import FrogWindow

class FrogApplication(Adw.Application):
    __gtype_name__ = 'FrogApplication'
    gtk_settings: Gtk.Settings

    settings: Settings = GObject.Property(type=GObject.TYPE_PYOBJECT)

    def __init__(self, version=None):
        super().__init__(application_id=APP_ID,
                         flags=Gio.ApplicationFlags.HANDLES_COMMAND_LINE)
        self.backend = None
        self.version = version

        # Init GSettings
        self.settings = Settings.new()

        self.add_main_option(
            'extract_to_clipboard',
            ord('e'),
            GLib.OptionFlags.NONE,
            GLib.OptionArg.NONE,
            _('Extract directly into the clipboard'),
            None
        )

        # Initialize tesseract data files storage.
        language_manager.init_tessdata()

        # Initialize libnotify.
        Notify.init("Frog")

    def do_startup(self, *args, **kwargs):
        Adw.Application.do_startup(self)

        # create command line option entries
        shortcut_entry = GLib.OptionEntry()
        shortcut_entry.long_name = 'extract_to_clipboard'
        shortcut_entry.short_name = ord('e')
        shortcut_entry.flags = 0
        shortcut_entry.arg = GLib.OptionArg.NONE
        shortcut_entry.arg_date = None
        shortcut_entry.description = _('Extract directly into the clipboard')
        shortcut_entry.arg_description = None

        self.backend = ScreenshotService()
        self.backend.connect('decoded', FrogApplication.on_decoded)

        action = Gio.SimpleAction.new("show_uri", GLib.VariantType.new('s'))
        action.connect("activate", self.on_show_uri)
        self.add_action(action)

        self.create_action('get_screenshot', self.get_screenshot, ['<primary>g'])
        self.create_action('get_screenshot_and_copy', self.get_screenshot_and_copy, ['<primary><shift>g'])
        self.create_action('copy_to_clipboard', self.on_copy_to_clipboard, ['<primary>g'])
        self.create_action('open_image', self.open_image, ['<primary>o'])
        self.create_action('paste_from_clipboard', self.on_paste_from_clipboard, ['<primary>v'])
        self.create_action('listen', self.on_listen, ['<primary>l'])
        self.create_action('listen_cancel', self.on_listen_cancel, ['<primary><ctrl>l'])
        self.create_action('shortcuts', self.on_shortcuts, ['<primary>question'])

        self.create_action('quit', lambda *_: self.quit(), ['<primary>q', '<primary>w'])
        self.create_action('about', self.on_about)
        self.create_action('preferences', self.on_preferences, ['<primary>comma'])

    def do_activate(self):
        win = self.props.active_window
        if not win:
            win = FrogWindow(application=self)
        win.present()

    def do_command_line(self, command_line):
        options = command_line.get_options_dict()
        options = options.end().unpack()

        if "extract_to_clipboard" in options:
            self.backend.capture(self.settings.get_string("active-language"), True)
            return 1

        self.activate()
        return 0

    def on_preferences(self, _action, _param) -> None:
        self.get_active_window().show_preferences()

    def on_about(self, _action, _param):
        about_window = Adw.AboutWindow(
            application_name="Frog",
            application_icon=APP_ID,
            version=self.version,
            copyright=f'© {datetime.date.today().year} Tender Owl',
            website="https://getfrog.app",
            issue_url="https://github.com/TenderOwl/Frog/issues/new",
            license_type=Gtk.License.MIT_X11,
            developer_name="TenderOwl Team",
            developers=["Andrey Maksimov"],
            release_notes="""
In this update, we've made a number of improvements to our app.

Share extracted text: This means you can now share your findings with the wider community.

Add Text-to-speech to be able to listen to the extracted text
                    without any additional apps (internet connection required).

Add the ability to insert images from the clipboard
                    by using the Ctrl+V shortcut.

In addition, we have upgraded the language selection dropdown
                    to popover with a quick search of a language.

Last but not least, Frog has become more accessibility-friendly.

We hope you enjoy our work!

            """,
            transient_for=self.props.active_window
        )
        about_window.present()

    def on_shortcuts(self, _action, _param):
        builder = Gtk.Builder()
        builder.add_from_resource(f"{RESOURCE_PREFIX}/ui/shortcuts.ui")
        builder.get_object("shortcuts").set_transient_for(self.get_active_window())
        builder.get_object("shortcuts").present()

    def on_copy_to_clipboard(self, _action, _param) -> None:
        self.get_active_window().on_copy_to_clipboard(self)

    def on_show_uri(self, _action, param) -> None:
        Gtk.show_uri(None, param.get_string(), Gdk.CURRENT_TIME)

    def get_screenshot(self, _action, _param) -> None:
        self.get_active_window().get_screenshot()

    def get_screenshot_and_copy(self, _action, _param) -> None:
        self.get_active_window().get_screenshot(copy=True)

    def open_image(self, _action, _param) -> None:
        self.get_active_window().open_image()

    def on_paste_from_clipboard(self, _action, _param) -> None:
        self.get_active_window().on_paste_from_clipboard(self)

    @staticmethod
    def on_decoded(_sender, text: str, copy: bool) -> None:
        icon = GdkPixbuf.Pixbuf.new_from_resource_at_scale(
            f"{RESOURCE_PREFIX}/icons/com.github.tenderowl.frog.svg",
            128, 128, True
        )

        if not text:
            notification: Notify.Notification = Notify.Notification.new(
                summary='Frog',
                body=_('No text found. Try to grab another region.')
            )
            notification.set_icon_from_pixbuf(icon)
            notification.show()

        if copy:
            clipboard_service.set(text)

            notification: Notify.Notification = Notify.Notification.new(
                summary='Frog',
                body=_('Text extracted. You can paste it with Ctrl+V')
            )
            notification.set_icon_from_pixbuf(icon)
            notification.show()

        else:
            print(f'{text}\n')

    def on_listen(self, _sender, _event):
        self.get_active_window().on_listen()

    def on_listen_cancel(self, _sender, _event):
        self.get_active_window().on_listen_cancel()

    def create_action(self, name, callback, shortcuts=None):
        """Add an application action.

        Args:
            name: the name of the action
            callback: the function to be called when the action is
              activated
            shortcuts: an optional list of accelerators
        """
        action = Gio.SimpleAction.new(name, None)
        action.connect("activate", callback)
        self.add_action(action)
        if shortcuts:
            self.set_accels_for_action(f"app.{name}", shortcuts)

def main(version):
    app = FrogApplication(version)
    return app.run(sys.argv)


""" Tests knockpy.knockoffs and knockpy.smatrix modules"""
import warnings
import numpy as np
import scipy as sp
import unittest

# for regular pytest calls
try:
    from .context import knockpy
# for running directly with python
except ImportError:
    from context import knockpy
from knockpy import dgp, utilities, mac, mrc, smatrix, knockoffs

try:
    import torch
    TORCH_AVAILABLE = True
except ImportError:
    TORCH_AVAILABLE = False
try:
    import choldate
    CHOLDATE_AVAILABLE = True
except:
    CHOLDATE_AVAILABLE = False
try:
    import pydsdp
    DSDP_AVAILABLE = True
except:
    DSDP_AVAILABLE = False

class CheckSMatrix(unittest.TestCase):

    # Helper function
    def check_S_properties(self, V, S, groups):

        # Test PSD-ness of S
        min_S_eig = np.linalg.eigh(S)[0].min()
        self.assertTrue(
            min_S_eig > 0,
            f"S matrix is not positive semidefinite: mineig is {min_S_eig}",
        )

        # Test PSD-ness of 2V - S
        min_diff_eig = np.linalg.eigh(2 * V - S)[0].min()
        self.assertTrue(
            min_diff_eig > 0,
            f"2Sigma-S matrix is not positive semidefinite: mineig is {min_diff_eig}",
        )

        # Calculate conditional knockoff matrix
        invV = utilities.chol2inv(V)
        invV_S = np.dot(invV, S)
        Vk = 2 * S - np.dot(S, invV_S)

        # Test PSD-ness of the conditional knockoff matrix
        min_Vk_eig = np.linalg.eigh(Vk)[0].min()
        self.assertTrue(
            min_Vk_eig > 0,
            f"conditional knockoff matrix is not positive semidefinite: mineig is {min_Vk_eig}",
        )

        # Test that S is just a block matrix
        p = V.shape[0]
        S_test = np.zeros((p, p))
        for j in np.unique(groups):

            # Select subset of S
            inds = np.where(groups == j)[0]
            full_inds = np.ix_(inds, inds)
            group_S = S[full_inds]

            # Fill only in this subset of S
            S_test[full_inds] = group_S

        # return
        np.testing.assert_almost_equal(
            S_test,
            S,
            decimal=5,
            err_msg="S matrix is not a block matrix of the correct shape",
        )

class TestEquicorrelated(CheckSMatrix):
    """ Tests equicorrelated knockoffs and related functions """

    def test_eigenvalue_calculation(self):

        # Test to make sure non-group and group versions agree
        # (in the case of no grouping)
        p = 100
        groups = np.arange(0, p, 1) + 1
        for rho in [0, 0.3, 0.5, 0.7]:
            V = np.zeros((p, p)) + rho
            for i in range(p):
                V[i, i] = 1
            expected_gamma = min(1, 2 * (1 - rho))
            gamma = mac.calc_min_group_eigenvalue(Sigma=V, groups=groups,)
            np.testing.assert_almost_equal(
                gamma,
                expected_gamma,
                decimal=3,
                err_msg="calc_min_group_eigenvalue calculates wrong eigenvalue",
            )

        # Test non equicorrelated version
        V = np.random.randn(p, p)
        V = np.dot(V.T, V) + 0.1 * np.eye(p)
        V = utilities.cov2corr(V)
        expected_gamma = min(1, 2 * np.linalg.eigh(V)[0].min())
        gamma = mac.calc_min_group_eigenvalue(Sigma=V, groups=groups)
        np.testing.assert_almost_equal(
            gamma,
            expected_gamma,
            decimal=3,
            err_msg="calc_min_group_eigenvalue calculates wrong eigenvalue",
        )

    def test_equicorrelated_construction(self):

        # Test S matrix construction
        p = 100
        groups = np.arange(0, p, 1) + 1
        V = np.random.randn(p, p)
        V = np.dot(V.T, V) + 0.1 * np.eye(p)
        V = utilities.cov2corr(V)

        # Expected construction
        expected_gamma = min(1, 2 * np.linalg.eigh(V)[0].min())
        expected_S = expected_gamma * np.eye(p)

        # Equicorrelated
        S = mac.solve_equicorrelated(Sigma=V, groups=groups)

        # Test to make sure the answer is expected
        np.testing.assert_almost_equal(
            S,
            expected_S,
            decimal=3,
            err_msg="calc_min_group_eigenvalue calculates wrong eigenvalue",
        )

    def test_psd(self):

        # Test S matrix construction
        p = 100
        V = np.random.randn(p, p)
        V = np.dot(V.T, V) + 0.1 * np.eye(p)
        V = utilities.cov2corr(V)

        # Create random groups
        groups = np.random.randint(1, p, size=(p))
        groups = utilities.preprocess_groups(groups)
        S = mac.solve_equicorrelated(Sigma=V, groups=groups)

        # Check S properties
        self.check_S_properties(V, S, groups)

class TestSDP(CheckSMatrix):
    """ Tests an easy case of SDP and ASDP """

    def test_easy_sdp(self):

        # Test non-group SDP first
        n = 200
        p = 50
        X, _, _, _, corr_matrix, groups = dgp.block_equi_graph(
            n=n, p=p, gamma=0.3
        )

        # S matrix
        trivial_groups = np.arange(0, p, 1) + 1
        S_triv = smatrix.compute_smatrix(
            Sigma=corr_matrix, groups=trivial_groups, method="sdp", verbose=True,
        )
        np.testing.assert_array_almost_equal(
            S_triv,
            np.eye(p),
            decimal=2,
            err_msg="solve_group_SDP does not produce optimal S matrix (blockequi graphs)",
        )
        self.check_S_properties(corr_matrix, S_triv, trivial_groups)

        # Repeat for gaussian_knockoffs method
        ksampler = knockoffs.GaussianSampler(
            X=X, Sigma=corr_matrix, groups=trivial_groups, verbose=False, method="sdp",
        )
        S_triv2 = ksampler.fetch_S()

        np.testing.assert_array_almost_equal(
            S_triv2,
            np.eye(p),
            decimal=2,
            err_msg="solve_group_SDP does not produce optimal S matrix (blockequi graphs)",
        )
        self.check_S_properties(corr_matrix, S_triv2, trivial_groups)

        # Test slightly harder case
        _, _, _, _, expected_out, _ = dgp.block_equi_graph(n=n, p=p, gamma=0)
        ksampler = knockoffs.GaussianSampler(
            X=X, Sigma=corr_matrix, groups=groups, verbose=False, method="sdp"
        )
        S_harder = ksampler.fetch_S()
        np.testing.assert_almost_equal(
            S_harder,
            expected_out,
            decimal=2,
            err_msg="solve_group_SDP does not produce optimal S matrix (blockequi graphs)",
        )
        self.check_S_properties(corr_matrix, S_harder, groups)

        # Repeat for ASDP
        ksampler = knockoffs.GaussianSampler(
            X=X,
            Sigma=corr_matrix,
            groups=groups,
            method="ASDP",
            verbose=False,
            max_block=10,
        )
        S_harder_ASDP = ksampler.fetch_S()
        np.testing.assert_almost_equal(
            S_harder_ASDP,
            expected_out,
            decimal=2,
            err_msg="solve_group_ASDP does not produce optimal S matrix (blockequi graphs)",
        )
        self.check_S_properties(corr_matrix, S_harder_ASDP, groups)

    def test_equicorr_SDP(self):

        # Test non-group SDP on equicorrelated cov matrix
        p = 100
        rhos = [0.6, 0.8, 0.9]
        solvers = ['sdp', 'cd']
        scales = [1, 5]
        for scale in scales:
            for rho in rhos:
                V = scale * rho * np.ones((p, p)) + scale * (1 - rho) * np.eye(p)
                for solver in solvers:
                    S = smatrix.compute_smatrix(Sigma=V, method="sdp", solver=solver, verbose=True)
                    expected = (2 - 2 * rho) * np.eye(p)
                    np.testing.assert_almost_equal(
                        S / scale,
                        expected,
                        decimal=2,
                        err_msg=f"compute_smatrix produces incorrect S_SDP for equicorr, rho={rho}, scale={scale}, solver={solver}",
                    )

    def test_cd_SDP(self):
        """
        Test that CD / SDP solvers get similar answers in a hard problem
        """

        np.random.seed(110)
        V = utilities.cov2corr(dgp.ErdosRenyi(p=100, tol=1e-1))
        S_cd = smatrix.compute_smatrix(
            Sigma=V, method="sdp", solver="cd", verbose=True, mu=0.95, num_iter=100
        )
        mac_cd = np.diag(S_cd).mean()
        S_sdp = smatrix.compute_smatrix(Sigma=V, method="sdp", solver="sdp")
        mac_sdp = np.diag(S_sdp).mean()
        np.testing.assert_almost_equal(
            mac_sdp,
            mac_cd,
            decimal=2,
            err_msg=f"compute_smatrix mac differs between SDP/CD solvers ({mac_sdp}, {mac_cd}, respectively"
        )

    def test_sdp_tolerance(self):

        # Get graph
        np.random.seed(110)
        Q = dgp.ErdosRenyi(p=50, tol=1e-1)
        V = utilities.cov2corr(utilities.chol2inv(Q))
        groups = np.concatenate([np.zeros(10) + j for j in range(5)]) + 1
        groups = groups.astype("int32")

        # Solve SDP
        for tol in [1e-3, 0.01, 0.02]:
            S = smatrix.compute_smatrix(
                Sigma=V,
                groups=groups,
                method="sdp",
                objective="pnorm",
                num_iter=10,
                tol=tol,
            )
            G = np.hstack([np.vstack([V, V - S]), np.vstack([V - S, V])])
            mineig = np.linalg.eig(G)[0].min()
            self.assertTrue(
                tol - mineig > -1 * tol / 10,
                f"sdp solver fails to control minimum eigenvalues: tol is {tol}, val is {mineig}",
            )
            self.check_S_properties(V, S, groups)

    def test_corrmatrix_errors(self):
        """ Tests that SDP raises informative errors when sigma is not scaled properly"""

        # Get graph
        np.random.seed(110)
        Q = dgp.ErdosRenyi(p=50, tol=1e-1)
        V = utilities.chol2inv(Q)
        groups = np.concatenate([np.zeros(10) + j for j in range(5)]) + 1
        groups = groups.astype("int32")

        # Helper function
        def SDP_solver():
            return mac.solve_group_SDP(V, groups)

        # Make sure the value error increases
        self.assertRaisesRegex(
            ValueError, "Sigma is not a correlation matrix", SDP_solver
        )

class TestCholupdate(CheckSMatrix):
    """
    Checks that the cholupdate function works correctly.
    """

    def test_cholupdate(self):

        # Generate complex covariance matrix and vector
        p = 500
        dgprocess = knockpy.dgp.DGP()
        dgprocess.sample_data(p=p, method='qer', delta=0.8)
        V = dgprocess.Sigma + np.eye(p)
        L = np.linalg.cholesky(V)
        np.testing.assert_array_almost_equal(
            V,
            np.dot(L, L.T),
            decimal=6,
            err_msg="numpy cholesky decomposition failed"
        )
        x = np.random.uniform(size=(p))

        # Test dense update
        V = V + np.outer(x, x)
        knockpy.mrc.cholupdate(L.T, x, add=True)
        np.testing.assert_array_almost_equal(
            V,
            np.dot(L, L.T),
            decimal=6,
            err_msg="mrc.cholupdate with add=True failed"
        )

        # Test sparse downdate
        x = np.zeros(p)
        x[5] = 0.2
        V = V - np.outer(x, x)
        knockpy.mrc.cholupdate(L.T, x, add=False)
        np.testing.assert_array_almost_equal(
            V,
            np.dot(L, L.T),
            decimal=6,
            err_msg="mrc.cholupdate with add=False failed"
        )

class TestDependencyWarnings(CheckSMatrix):
    """ Tests DSDP_warning / choldate_warning """

    def test_warnings_raised(self):

        print(f"torch={TORCH_AVAILABLE}, dsdp={DSDP_AVAILABLE}, choldate={CHOLDATE_AVAILABLE}")

        # Sigma
        p = 50
        rho = 0.8
        Sigma = rho*np.ones((p,p)) + (1-rho)*np.eye(p)

        for method, dependency_flag, dependency_name, suppress_kwargs in zip(
            ['mvr', 'maxent', 'sdp'],
            [CHOLDATE_AVAILABLE, CHOLDATE_AVAILABLE, DSDP_AVAILABLE],
            ["choldate", "choldate", "dsdp"],
            [{"choldate_warning":False}, {"choldate_warning":False}, {"dsdp_warning":False}]
        ):
            # Check default setting
            with warnings.catch_warnings(record=True) as w:
                warnings.simplefilter('always')
                smatrix.compute_smatrix(Sigma, method=method)
                if not dependency_flag:
                    self.assertTrue(
                        dependency_name in str(w[-1].message)
                    )
                else:
                    self.assertTrue(len(w) == 0)
            # Check that it does not trigger with correct flag
            with warnings.catch_warnings(record=True) as w:
                warnings.simplefilter('always')
                smatrix.compute_smatrix(Sigma, method=method, **suppress_kwargs)
                self.assertTrue(len(w) == 0, "warning still triggers with choldate_warning=False")

class TestMRCSolvers(CheckSMatrix):
    """ Tests the various MRC solvers / classes"""

    def test_scale_sqrt_S(self):
        """ Tests the function which scales sqrt S for PSGD solver"""

        if not TORCH_AVAILABLE:
            return None
        from knockpy import kpytorch

        # Construct covariance matrix
        p = 50
        rho = 0.8
        Sigma = np.zeros((p, p)) + rho
        Sigma += (1 - rho) * np.eye(p)
        groups = np.arange(1, p + 1, 1)
        init_blocks = [np.eye(1) for _ in range(p)]

        # Create model - this automatically scales the
        # initial blocks properly
        fk_precision_calc = kpytorch.mrcgrad.MVRLoss(Sigma, groups, init_S=init_blocks)
        # Check for proper scaling
        S = fk_precision_calc.pull_S().detach().numpy()
        expected = min(1, 2 - 2 * rho) * np.eye(p)
        np.testing.assert_almost_equal(
            S,
            expected,
            decimal=1,
            err_msg=f"Initial scaling fails, expected {expected} but got {S} for equicorrelated rho={rho}",
        )

    def test_group_sorting_error(self):
        """ Tests PSGD class raises error if the cov
        matrix/groups are improperly sorted"""

        if not TORCH_AVAILABLE:
            return None
        from knockpy import kpytorch

        # Groups and sigma
        p = 50
        Sigma = np.eye(p)
        groups = knockpy.utilities.preprocess_groups(np.random.randint(1, p + 1, p))

        # Try to initialize
        def init_unsorted_model():
            model = kpytorch.mrcgrad.MVRLoss(Sigma, groups)

        self.assertRaisesRegex(
            ValueError,
            "Sigma and groups must be sorted prior to input",
            init_unsorted_model,
        )

    def test_smoothing(self):
        """
        Smoothing is not required for this, but this is a nice check anyway.
        """
        p = 50
        smoothing = 0.1
        dgprocess = dgp.DGP()
        _, _, _, _, V = dgprocess.sample_data(method="partialcorr", rho=0.1,)
        S_MVR = mrc.solve_mvr(Sigma=V, smoothing=smoothing)
        # Not implemented yet
        S_SDP = mac.solve_SDP(Sigma=V, tol=1e-5)
        mvr_mean = np.diag(S_MVR).mean()
        sdp_mean = np.diag(S_SDP).mean()
        self.assertTrue(
            sdp_mean - mvr_mean < 1e-3,
            f"Highly smoothed S_MVR ({S_MVR}) too far from S_SDP ({S_SDP}) for equi partial corr",
        )

    def test_equicorrelated_soln(self):
        """ Tests that solvers yield expected
        solution for equicorrelated matrices """

        if not TORCH_AVAILABLE:
            return None
        from knockpy import kpytorch

        # Main constants
        p = 50
        groups = np.arange(1, p + 1, 1)
        smoothings = [0, 0.01]

        # Construct equicorrelated matrices
        rhos = [0.1, 0.3, 0.5, 0.7, 0.9]
        for rho in rhos:
            for smoothing in smoothings:
                for method in ["mvr", "mmi", "maxent"]:
                    # Construct Sigma
                    Sigma = np.zeros((p, p)) + rho
                    Sigma = Sigma + (1 - rho) * np.eye(p)

                    # Expected solution
                    opt_prop_rec = min(rho, 0.5)
                    max_S_val = min(1, 2 - 2 * rho)
                    expected = (1 - rho) * np.eye(p)

                    # Test optimizer
                    opt_S = kpytorch.mrcgrad.solve_mrc_psgd(
                        Sigma=Sigma,
                        groups=groups,
                        init_S=None,
                        smoothing=smoothing,
                        verbose=True,
                        tol=1e-5,
                        max_epochs=100,
                        line_search_iter=10,
                        lr=1e-2,
                    )
                    self.check_S_properties(Sigma, opt_S, groups)
                    np.testing.assert_almost_equal(
                        opt_S,
                        expected,
                        decimal=1,
                        err_msg=f"For equicorrelated cov rho={rho}, PSGD solver yields unexpected solution",
                    )

                # Test MVR coordinate descent optimizer
                opt_S = mrc.solve_mvr(Sigma=Sigma, smoothing=smoothing, verbose=True)
                self.check_S_properties(Sigma, opt_S, groups)
                np.testing.assert_almost_equal(
                    opt_S,
                    expected,
                    decimal=2,
                    err_msg=f"For equicorrelated cov rho={rho}, mvr_solver yields unexpected solution",
                )

                # Test maximum entropy coordinate descent optimizer
                if smoothing == 0:
                    opt_S = mrc.solve_mmi(
                        Sigma=Sigma, verbose=True
                    )
                    self.check_S_properties(Sigma, opt_S, groups)
                    np.testing.assert_almost_equal(
                        opt_S,
                        expected,
                        decimal=2,
                        err_msg=f"For equicorrelated cov rho={rho}, mmi_solver yields unexpected solution",
                    )

    def test_maxent(self):
        """ Both maxent/mmi work properly """
        # Sample data
        dgprocess = dgp.DGP()
        dgprocess.sample_data(p=50, method='ar1', a=3)

        # Check solve_maxent/solve_mmi
        np.random.seed(110)
        S_ME = smatrix.compute_smatrix(dgprocess.Sigma, method='maxent')
        np.random.seed(110)
        S_MMI = smatrix.compute_smatrix(dgprocess.Sigma, method='mmi')
        np.testing.assert_array_almost_equal(
            S_ME, S_MMI, decimal=3, err_msg=f"compute_smatrix yields diff answers for mmi, maxent"
        )

        # Check solve_maxent/solve_mmi
        np.random.seed(110)
        S_ME = mrc.solve_maxent(dgprocess.Sigma)
        np.random.seed(110)
        S_MMI = mrc.solve_mmi(dgprocess.Sigma)
        np.testing.assert_array_almost_equal(
            S_ME, S_MMI, decimal=3, err_msg=f"solve_maxent and solve_mmi yield different answers"
        )

        # Check maxent_loss/mmi_loss
        L_ME = mrc.maxent_loss(dgprocess.Sigma, S_ME)
        L_MMI = mrc.mmi_loss(dgprocess.Sigma, S_MMI)
        np.testing.assert_almost_equal(
            L_ME, L_MMI, decimal=3, err_msg=f"maxent_loss and mmi_loss yield different answers"
        )

    def test_equicorrelated_soln_recycled(self):
        """ Correct solutions on equicorrelated matrix with recycling """

        # Potentially also test the PSGD solver if torch is available
        if TORCH_AVAILABLE:
            from knockpy import kpytorch

        # Main constants
        p = 50
        groups = np.arange(1, p + 1, 1)

        # Test separately with the recycling proportion param
        rhos = [0.1, 0.3, 0.5, 0.7, 0.9]
        true_rec_props = [0.5, 0.25, 0.8, 0.5, 0.5]
        for true_rec_prop, rho in zip(true_rec_props, rhos):

            # Construct Sigma
            Sigma = np.zeros((p, p)) + rho
            Sigma = Sigma + (1 - rho) * np.eye(p)

            # Expected solution
            opt_prop_rec = min(rho, 0.5)
            max_S_val = min(1, 2 - 2 * rho)
            normal_opt = (1 - opt_prop_rec) * max_S_val
            new_opt = min(2 - 2 * rho, normal_opt / (1 - true_rec_prop))
            expected = new_opt * np.eye(p)

            # Test PSGD optimizer
            if TORCH_AVAILABLE:
                opt_S = kpytorch.mrcgrad.solve_mrc_psgd(
                    Sigma=Sigma,
                    groups=groups,
                    init_S=None,
                    rec_prop=true_rec_prop,
                    tol=1e-5,
                    max_epochs=100,
                    line_search_iter=10,
                )
                self.check_S_properties(Sigma, opt_S, groups)
                np.testing.assert_almost_equal(
                    opt_S,
                    expected,
                    decimal=2,
                    err_msg=f"For equicorrelated cov rho={rho} rec_prop={true_rec_prop}, MVR PSDSolver returns {opt_S}, expected {expected}",
                )

            # Coord descent solver
            np.random.seed(110)
            opt_S = mrc.solve_mvr(
                Sigma=Sigma, rej_rate=true_rec_prop, tol=1e-5, verbose=True,
            )
            self.check_S_properties(Sigma, opt_S, groups)
            np.testing.assert_almost_equal(
                opt_S,
                expected,
                decimal=2,
                err_msg=f"For equicorrelated cov rho={rho} rec_prop={true_rec_prop}, MVR CD solver returns {opt_S}, expected {expected}",
            )

    def test_complex_solns(self):
        """
        Check the solution of the various solvers
        for non-grouped knockoffs.
        """

        # Check availability
        if not TORCH_AVAILABLE:
            return None
        from knockpy import kpytorch

        np.random.seed(110)
        p = 100
        methods = ["ar1", "ver"]
        groups = np.arange(1, p + 1, 1)
        for method in methods:
            dgprocess = dgp.DGP()
            _, _, _, _, Sigma = dgprocess.sample_data(method=method, p=p)

            # Use SDP as baseline
            init_S = knockpy.mac.solve_group_SDP(Sigma, groups)
            sdp_mvr_loss = mrc.mvr_loss(Sigma, init_S)

            # Apply gradient solver
            opt_S = kpytorch.mrcgrad.solve_mrc_psgd(
                Sigma=Sigma,
                groups=groups,
                init_S=init_S,
                tol=1e-5,
                max_epochs=100,
                line_search_iter=10,
            )
            psgd_mvr_loss = mrc.mvr_loss(Sigma, opt_S)

            # Check S matrix
            self.check_S_properties(Sigma, opt_S, groups)
            # Check new loss < init_loss
            self.assertTrue(
                psgd_mvr_loss <= sdp_mvr_loss,
                msg=f"For {method}, PSGD solver has higher loss {psgd_mvr_loss} v. sdp {sdp_mvr_loss}",
            )

            # MVR solver outperforms PSGD
            opt_S_mvr = mrc.solve_mvr(Sigma=Sigma)
            self.check_S_properties(Sigma, opt_S_mvr, groups)
            cd_mvr_loss = mrc.mvr_loss(Sigma, opt_S_mvr)
            self.assertTrue(
                cd_mvr_loss <= psgd_mvr_loss,
                msg=f"For {method}, coord descent MVR solver has higher loss {cd_mvr_loss} v. PSGD {psgd_mvr_loss}",
            )

            # mmi solver outperforms PSGD
            print(Sigma)
            print(np.linalg.eigh(Sigma)[0].min())
            opt_S_mmi = mrc.solve_mmi(Sigma=Sigma)
            print(opt_S_mmi)
            print(opt_S)
            self.check_S_properties(Sigma, opt_S_mmi, groups)
            cd_mmi_loss = mrc.mmi_loss(Sigma, opt_S_mmi)
            psgd_mmi_loss = mrc.mmi_loss(Sigma, opt_S)
            self.assertTrue(
                cd_mmi_loss <= psgd_mmi_loss,
                msg=f"For {method}, coord descent mmi solver has higher loss {cd_mmi_loss} v. PSGD {psgd_mmi_loss}",
            )

    def test_cd_group_mvr(self):

        np.random.seed(110)
        p = 100
        for method in ['ar1', 'ver', 'qer']:
            # Sample data
            dgprocess = knockpy.dgp.DGP()
            dgprocess.sample_data(method=method, p=p)
            Sigma = dgprocess.Sigma

            # Sanity check: grouped solver gets right soln for ungrouped
            S1 = mrc._solve_mvr_grouped(
                Sigma, groups=np.arange(1, p+1), converge_tol=1e-2, verbose=True
            )
            S2 = mrc.solve_mvr(Sigma)
            Sdiff = S1 - S2
            mean_diff = np.abs(Sdiff[Sdiff != 0]).mean()
            err_msg = f"For no groups, method={method}, avg diff. between grouped/ungrouped solvers is {mean_diff} > 1e-2"
            err_msg = f"grouped solver: {S1} \n ungrouped solver: {S2} \n" + err_msg
            self.assertTrue(
                mean_diff < 1e-2, err_msg
            )

            # Check that adding groups lowers loss
            loss_ungrouped = mrc.mvr_loss(Sigma, S2)
            groups = np.around(np.arange(1, p+1)/2)
            Sgroup = mrc.solve_mvr(Sigma, groups)
            loss_grouped = mrc.mvr_loss(Sigma, Sgroup)
            self.assertTrue(
                loss_grouped < loss_ungrouped,
                f"Grouped loss ({loss_grouped}) > ungrouped loss ({loss_ungrouped}) for method={method}"
            )

    def test_complex_group_solns(self):
        """
        Check the solutions of the PSGD solver
        for group knockoffs.
        """

        if not TORCH_AVAILABLE:
            return None
        from knockpy import kpytorch

        # Construct graph + groups
        np.random.seed(110)
        p = 50
        groups = knockpy.utilities.preprocess_groups(np.random.randint(1, p + 1, p))
        for method in ["ar1", "ver"]:
            dgprocess = dgp.DGP()
            _, _, _, _, Sigma = dgprocess.sample_data(method=method, p=p,)

            # Use SDP as baseline
            init_S = knockpy.mac.solve_group_SDP(Sigma, groups)
            init_loss = mrc.mvr_loss(Sigma, init_S)

            # Apply gradient solver
            opt_S = kpytorch.mrcgrad.solve_mrc_psgd(
                Sigma=Sigma,
                groups=groups,
                init_S=init_S,
                tol=1e-5,
                max_epochs=100,
                line_search_iter=10,
            )
            psgd_loss = mrc.mvr_loss(Sigma, opt_S)

            # Check S matrix
            self.check_S_properties(Sigma, opt_S, groups)
            # Check new loss < init_loss
            self.assertTrue(
                psgd_loss <= init_loss,
                msg=f"For {method}, PSGD solver has higher loss {psgd_loss} v. sdp {init_loss}",
            )

    def test_equi_ciknock_solution(self):
        """
        Check ciknockoff solution
        """
        p = 500
        rho = 0.6
        # 1. Block equicorrelated
        dgprocess = dgp.DGP()
        _, _, _, _, Vblock = dgprocess.sample_data(
            p=p, method="blockequi", gamma=0, rho=rho, block_size=2
        )
        S_CI = mrc.solve_ciknock(Vblock)
        np.testing.assert_almost_equal(
            S_CI,
            (1 - rho ** 2) * np.eye(p),
            2,
            "S_CI is incorrect for block-equicorrelated with blocksize 2",
        )
        # 2. Equicorelated
        dgprocess = dgp.DGP()
        _, _, _, _, V = dgprocess.sample_data(
            p=p, method="blockequi", gamma=1, rho=rho
        )
        S_CI = mrc.solve_ciknock(V)
        np.testing.assert_almost_equal(
            S_CI, (1 - rho) * np.eye(p), 2, "S_CI is incorrect for equicorrelated"
        )

class TestBlockdiagApprx(CheckSMatrix):
    def test_divide_comp(self):
        """ Check that when we divide computation, our
        final blocks do not exceed the specified block size."""

        p = 120
        for method in ["blockequi", "ar1", "ver", "qer"]:
            dgprocess = dgp.DGP()
            _, _, _, _, V = dgprocess.sample_data(p=p, method=method)
            for max_block in [3, 12, 15, 54, 76]:
                blocks = smatrix.divide_computation(V, max_block)
                block_sizes = utilities.calc_group_sizes(blocks)
                self.assertTrue(
                    block_sizes.max() <= max_block,
                    msg=f"Returned blocks have max size {block_sizes.max()} > max_block {max_block}",
                )

    def test_blockdiag_approx_on_blockdiag(self):

        # Check that blockdiag approx yields same solution
        # for block-diagonal matrix.
        p = 250
        rho = 0.3
        gamma = 0
        groups = np.arange(1, p + 1, 1)
        dgprocess = dgp.DGP()
        _, _, _, _, V = dgprocess.sample_data(
            p=p, method="blockequi", gamma=gamma, rho=rho
        )

        for method in ["mvr", "maxent", "sdp"]:
            # Check that S_approx is valid
            S_approx = smatrix.compute_smatrix(V, method=method, max_block=10)
            self.check_S_properties(V, S_approx, groups)

            # In this case, should be the same as S_exact
            S_exact = smatrix.compute_smatrix(V, method=method)
            np.testing.assert_array_almost_equal(
                S_approx,
                S_exact,
                2,
                err_msg="Blockdiag approximation yields incorrect answer for blockdiag Sigma",
            )

    def test_blockdiag_approx_on_ar1(self):

        # Check that blockdiag approx yields a good soln on AR1
        p = 150
        a = 1
        b = 1
        # Create trivial and nontrivial groups
        triv_groups = np.arange(1, p + 1, 1)
        nontriv_groups = np.around(triv_groups / 2 + 0.01)
        nontriv_groups = utilities.preprocess_groups(nontriv_groups)
        # Sample data
        np.random.seed(110)
        dgprocess = dgp.DGP()
        _, _, _, _, V = dgprocess.sample_data(p=p, method="ar1", a=a, b=b)
        for groups in [triv_groups, nontriv_groups]:
            for method in ["mvr", "maxent", "mmi", "sdp"]:

                # Note we can't fit group mvr/maxent without pytorch
                nontriv_group_flag = np.all(groups == nontriv_groups)
                if not TORCH_AVAILABLE and nontriv_group_flag:
                    continue

                # Check that S_approx is valid
                S_approx = smatrix.compute_smatrix(
                    V, method=method, groups=groups, max_block=50
                )
                self.check_S_properties(V, S_approx, groups)

                # For exponentially decaying offdiags, should be quite similar
                # This seems to not be true for groups, so skip that for now...
                # For now, skip grouped mvr/mmi
                if nontriv_group_flag:
                    continue
                S_exact = smatrix.compute_smatrix(V, method=method, groups=groups)
                diff = np.abs(S_approx - S_exact)
                mean_diff = diff[S_exact != 0].mean()
                expected = 1e-2 if CHOLDATE_AVAILABLE else 5e-2
                identifier = f"for method={method}, groups={groups}"
                self.assertTrue(
                    mean_diff < expected,
                    msg=f"Blockdiag apprx is {mean_diff} > {expected} on avg. away from exact soln {identifier}",
                )

class TestFactorApprx(CheckSMatrix):

    def test_agreement_on_factor_models(self):

        # Random factor models
        np.random.seed(110)
        ks = [1, 10, 20, 50]
        p = 150
        for k in ks:
            for method in ['mvr', 'maxent', 'sdp']:
                D = np.random.uniform(low=0.01, high=1, size=(p,))
                U = np.random.randn(p, k) / np.sqrt(k)
                # Rescale to correlation matrix
                diag_Sigma = D + (np.power(U, 2)).sum(axis=1)
                D = D / diag_Sigma
                U = U / np.sqrt(diag_Sigma).reshape(-1, 1)
                Sigma = np.diag(D) + np.dot(U, U.T)
                # Check the factor method agrees with reg. method
                S1 = smatrix.compute_smatrix(
                    Sigma=None, D=D, U=U, method=method, how_approx='factor'
                )
                S2 = smatrix.compute_smatrix(
                   Sigma=Sigma, method=method, solver='cd', tol=0
                )
                np.testing.assert_array_almost_equal(
                    S1, S2, decimal=2,
                    err_msg=f"factored/non-factored versions for method={method}, k={k} do not agree"
                )

class CheckValidKnockoffs(unittest.TestCase):
    def check_valid_mxknockoffs(self, X, mu=None, Sigma=None, msg="", **kwargs):

        # S matrix
        ksampler = knockoffs.GaussianSampler(
            X=X, mu=mu, Sigma=Sigma, verbose=True, **kwargs
        )
        S = ksampler.fetch_S()
        Xk = ksampler.sample_knockoffs()

        # Test knockoff mean
        if mu is None:
            mu = X.mean(axis=0)
        outmsg = "Knockoffs have incorrect means "
        outmsg += f"for MX knockoffs for {msg}"
        np.testing.assert_array_almost_equal(Xk.mean(axis=0), mu, 2, outmsg)

        # Sigma should be
        if Sigma is None:
            Sigma, _ = utilities.estimate_covariance(X, tol=1e-2)

        # Also rescale X/Xk so S makes sense
        scale = np.sqrt(np.diag(Sigma))
        X = X / scale.reshape(1, -1)
        if mu is None:
            mu = X.mean(axis=0)
        else:
            mu = mu / scale
        Xk = Xk / scale.reshape(1, -1)
        Sigma = Sigma / np.outer(scale, scale)
        S = S / np.outer(scale, scale)

        # Empirical FK correlation matrix
        features = np.concatenate([X, Xk], axis=1)
        G_hat = np.cov(features.T)

        # Calculate population version
        G = np.concatenate(
            [np.concatenate([Sigma, Sigma - S]), np.concatenate([Sigma - S, Sigma])],
            axis=1,
        )

        # Test G has correct structure
        outmsg = f"Feature-knockoff cov matrix has incorrect values "
        outmsg += f"for MX knockoffs for {msg} graph "
        np.testing.assert_array_almost_equal(G_hat, G, 2, outmsg)

class TestKnockoffGen(CheckValidKnockoffs):
    """ Tests whether knockoffs have correct distribution empirically"""

    def test_method_parser(self):

        # Easiest test
        method1 = "hello"
        out1 = smatrix.parse_method(method1, None, None)
        self.assertTrue(
            out1 == method1, "parse method fails to return non-None methods"
        )

        # Default is mvr
        p = 1000
        groups = np.arange(1, p + 1, 1)
        out2 = smatrix.parse_method(None, groups, p)
        self.assertTrue(out2 == "mvr", "parse method fails to return mvr by default")

        # Otherwise SDP
        groups[-1] = 1
        out2 = smatrix.parse_method(None, groups, p)
        self.assertTrue(
            out2 == "sdp", "parse method fails to return SDP for grouped knockoffs"
        )

    def test_error_raising(self):

        # Generate data
        n = 100
        p = 100
        X, _, _, _, corr_matrix, groups = dgp.block_equi_graph(
            n=n, p=p, gamma=1, rho=0.8
        )
        S_bad = np.eye(p)

        def fdr_vio_knockoffs():
            ksampler = knockoffs.GaussianSampler(
                X=X, Sigma=corr_matrix, S=S_bad, verbose=False
            )
            ksampler.sample_knockoffs(check_psd=True)

        self.assertRaisesRegex(
            np.linalg.LinAlgError,
            "meaning FDR control violations are extremely likely",
            fdr_vio_knockoffs,
        )

        # Test FX knockoff violations
        def fx_knockoffs_low_n():
            knockoffs.FXSampler(
                X=X, Sigma=corr_matrix, S=None,
            )

        self.assertRaisesRegex(
            np.linalg.LinAlgError,
            "FX knockoffs can't be generated with n",
            fx_knockoffs_low_n,
        )

    def test_consistency_of_inferring_sigma(self):
        """ Checks that the same knockoffs are produced
        whether you infer the covariance matrix first and
        pass it to the gaussian_knockoffs generator, or
        you let the generator do the work for you
        """

        n = 25
        p = 300
        rho = 0.5
        dgprocess = dgp.DGP()
        X, _, _, _, _ = dgprocess.sample_data(n=n, p=p, rho=rho, method="AR1")

        # Method 1: infer cov first
        V, _ = utilities.estimate_covariance(X, tol=1e-2)
        np.random.seed(110)
        Xk1 = knockoffs.GaussianSampler(
            X=X, Sigma=V, method="sdp"
        ).sample_knockoffs()

        # Method 2: Infer during
        np.random.seed(110)
        Xk2 = knockoffs.GaussianSampler(
            X=X, method="sdp"
        ).sample_knockoffs()
        np.testing.assert_array_almost_equal(
            Xk1, Xk2, 5, err_msg="Knockoff gen is inconsistent"
        )

    def test_MX_knockoff_dist(self):

        # Test knockoff construction for mvr and SDP
        # on equicorrelated matrices
        np.random.seed(110)
        n = 100000
        copies = 3
        p = 5

        # Check with a non-correlation matrix
        V = 4 * dgp.AR1(p=p, rho=0.5)
        mu = np.random.randn(p)
        print(f"true mu: {mu}")
        dgprocess = dgp.DGP(mu=mu, Sigma=V)
        X, _, _, _, _ = dgprocess.sample_data(n=n, p=p,)
        print(f"X mean: {X.mean(axis=0)}")

        # Check validity for oracle cov matrix
        self.check_valid_mxknockoffs(X, mu=mu, Sigma=V, msg=f"ORACLE 4*AR1(rho=0.5)")

        # Check validity for estimated cov matrix
        print("Here now")
        self.check_valid_mxknockoffs(X, msg=f"ESTIMATED 4*AR1(rho=0.5)")

        # Check for many types of data
        for rho in [0.1, 0.9]:
            for gamma in [0.5, 1]:
                for method in ["mvr", "sdp"]:

                    mu = 10 * np.random.randn(p)
                    X, _, _, _, corr_matrix, _ = dgp.block_equi_graph(
                        n=n, p=p, gamma=gamma, rho=rho, mu=mu
                    )

                    # Check validity for oracle correlation matrix
                    self.check_valid_mxknockoffs(
                        X,
                        mu=mu,
                        Sigma=corr_matrix,
                        msg=f"blockequi graph, rho = {rho}, gamma = {gamma}",
                    )

                    # Check validity for estimation
                    self.check_valid_mxknockoffs(
                        X,
                        msg=f"ESTIMATED blockequi graph, rho = {rho}, gamma = {gamma}",
                    )

    def test_FX_knockoff_dist(self):
        # Test knockoff construction for mvr and SDP
        # on equicorrelated matrices
        n = 500
        p = 5
        for rho in [0.1, 0.9]:
            for gamma in [0.5, 1]:
                for method in ["mvr", "sdp"]:
                    # X values
                    X, _, _, _, corr_matrix, _ = dgp.block_equi_graph(
                        n=n, p=p, gamma=gamma, rho=rho
                    )
                    # S matrix
                    trivial_groups = np.arange(0, p, 1) + 1
                    ksampler = knockoffs.FXSampler(X=X, method=method, verbose=False)
                    Xk = ksampler.sample_knockoffs()
                    S = ksampler.fetch_S()

                    # Compute empirical (scaled) cov matrix
                    features = np.concatenate([X, Xk], axis=1)
                    G_hat = np.dot(features.T, features)

                    # Calculate what this should be
                    Sigma = np.dot(X.T, X)
                    G = np.concatenate(
                        [
                            np.concatenate([Sigma, Sigma - S]),
                            np.concatenate([Sigma - S, Sigma]),
                        ],
                        axis=1,
                    )

                    # Test G has correct structure
                    msg = f"Feature-knockoff cov matrix has incorrect values "
                    msg += f"for blockequi graph, FX knockoffs, rho = {rho}, gamma = {gamma}"
                    np.testing.assert_array_almost_equal(G_hat, G, 5, msg)

    def test_scaling_consistency(self):
        """
        Checks whether if we first calculate S
        and then pass it back into the knockoff
        generator, we'll get the same answer back.
        """

        p = 100
        n = 300
        dgprocess = dgp.DGP()
        X, y, beta, Q, V = dgprocess.sample_data(
            method="qer", p=p, n=n, coeff_size=0.5, sparsity=0.5,
        )
        ksampler1 = knockoffs.FXSampler(X=X,)
        S1 = ksampler1.fetch_S()
        ksampler2 = knockoffs.FXSampler(X=X, S=S1,)
        S2 = ksampler2.fetch_S()
        np.testing.assert_array_almost_equal(
            S1,
            S2,
            decimal=6,
            err_msg=f"Repeatedly passing S into/out of knockoff gen yields inconsistencies",
        )

if __name__ == "__main__":
    import pytest
    import sys
    pytest.main(sys.argv)
    #unittest.main()