# Chapter 5: Data Displays and Grids
# Recipe 6: Modelling your data
#
import inspect
import wx
import wx.dataview as dv

class HierarchyInfo:
    def __init__(self, item, parent):
        self.item = item
        self.parent = parent

        self.name = item.__name__
        self.docs = item.__doc__
        if self.docs:
            self.docs = self.docs.replace("\n", "")
        self.subs = list()

        self._searchSubs()

    def _searchSubs(self):
        if hasattr(self.item, '__subclasses__'):
            for t in self.item.__subclasses__():
                self.subs.append(HierarchyInfo(t, self))

class ClassDataModel(dv.PyDataViewModel):
    def __init__(self, data):
        super(ClassDataModel, self).__init__()
        self.data = data
        self.objmapper.UseWeakRefs(True)

    def GetChildren(self, parent, children):
        # check root node
        if not parent:
            for item in self.data:
                children.append(self.ObjectToItem(item))
            return len(self.data)

        node = self.ItemToObject(parent)
        for item in node.subs:
            children.append(self.ObjectToItem(item))
        return len(node.subs)

    def HasContainerColumns(self, item):
        return True

    def IsContainer(self, item):
        if not item:
            return True

        obj = self.ItemToObject(item)
        return len(obj.subs) > 0

    def GetParent(self, item):
        if not item:
            return dv.NullDataViewItem

        obj = self.ItemToObject(item)
        if obj.parent is None:
            return dv.NullDataViewItem
        else:
            return self.ObjectToItem(obj.parent)

    def GetValue(self, item, col):
        obj = self.ItemToObject(item)
        vMap = { 0 : obj.name,
                 1 : len(obj.subs),
                 2 : obj.docs
               }
        return vMap[col]

    def SetValue(self, value, item, col):
        pass

    def GetAttr(self, item, col, attr):
        if col == 1:
            attr.Bold = True
            return True
        return False

#------- Sample Application ---------#

class ClassViewer(wx.Frame):
    def __init__(self, parent, title):
        super(ClassViewer, self).__init__(parent, title=title)

        # Look at all classes in wx namespace
        data = list()
        for x in dir(wx):
            item = getattr(wx, x)
            if inspect.isclass(item):
                data.append(HierarchyInfo(item, None))

        dvc = dv.DataViewCtrl(self, style=dv.DV_VERT_RULES)
        model = ClassDataModel(data)
        dvc.AssociateModel(model)

        autosize = wx.COL_WIDTH_AUTOSIZE
        dvc.AppendTextColumn("Class", 0, width=autosize)
        dvc.AppendTextColumn("Subclasses", 1, width=autosize,
                             align=wx.ALIGN_CENTER)
        dvc.AppendTextColumn("Docstring", 2, width=autosize)

        sizer = wx.BoxSizer()
        sizer.Add(dvc, 1, wx.EXPAND)
        self.SetSizer(sizer)
        self.SetInitialSize((400, 200))

class MyApp(wx.App):
    def OnInit(self):
        self.frame = ClassViewer(None, title="DataView Model")
        self.frame.Show();
        return True

if __name__ == "__main__":
    app = MyApp(False)
    app.MainLoop()


import streamlit as st
from streamlit import session_state as ss
import pandas as pd
import numpy as np
from shapely import Polygon
import shapely

from bokeh.plotting import figure
from bokeh.models import ColumnDataSource, FixedTicker, HoverTool, ColorBar
from bokeh.palettes import Inferno, Inferno256, Magma, Magma256, Greys, Greys256, Plasma, Plasma256, Viridis, Viridis256, Cividis, Cividis256, Turbo, Turbo256
from bokeh.transform import linear_cmap, log_cmap
from bokeh.embed import file_html
from bokeh.resources import CDN
from bokeh.models import CustomJSHover, HoverTool

import math

def read_file(file: str) -> pd.DataFrame:
    if '.xlsx' in file.name:
        df = pd.read_excel(file, sheet_name=0)
    elif '.csv' in file.name:
        df = pd.read_csv(file)
    else:
        df = None

    return df

def get_plot_columns(df: pd.DataFrame) -> (str, str):
    # st.write(df.dtypes)
    plot_columns = df.select_dtypes(['int', 'float']).columns
    x_column = st.selectbox(
        label='Select X Column for plot',
        options=plot_columns
    )
    y_column = st.selectbox(
        label='Select Y Column for plot',
        options=plot_columns
    )

    return x_column, y_column

def get_additional_column(df: pd.DataFrame) -> (str, str):
    plot_columns = df.select_dtypes(['int', 'float']).columns

    z_column = st.selectbox(
        label='Select Z Column for plot',
        options=plot_columns
    )

    palette = st.selectbox(
        label='Select a color palette',
        options=[
            'Cividis',
            'Inferno',
            'Magma',
            'Plasma',
            'Viridis',
            'Turbo',
            'Grey'
        ]
    )

    return z_column, palette

def plot(df: pd.DataFrame, x_column: str, y_column: str, ratio: float, holes: bool) -> None:
    points = df[[x_column, y_column]].to_numpy()
    coords = list(map(tuple, np.asarray(points)))
    poly = Polygon(coords)
    x, y = shapely.concave_hull(poly, ratio=ratio, allow_holes=holes).exterior.xy

    p = figure(
		x_axis_label=x_column,
		y_axis_label=y_column,
		aspect_ratio = 17.5 / 8,
		sizing_mode='scale_width',
		output_backend='webgl'
	)
    source = ColumnDataSource(df)
    p.circle(
        source=source,
        x=x_column,
        y=y_column
    )
    p.line(
        x=x,
        y=y,
        color='black'
    )
    p.toolbar.logo = None

    st.bokeh_chart(p)

    hull_df = pd.DataFrame({'Hull X Coordinates' : x, 'Hull Y Coordinates': y})
    st.download_button(
        label='Download Hull Coordinates',
        data=hull_df.to_csv(index=False).encode('utf-8'),
        file_name='Hull Coordinates.csv',
        mime='text/csv'
    )

def plot_contour(df: pd.DataFrame, x_column: str, y_column: str, z_column: str, ratio: float, holes: bool, palette: list, log_scale: bool, small_palette_num: int) -> None:
    points = df[[x_column, y_column]].to_numpy()
    coords = list(map(tuple, np.asarray(points)))
    poly = Polygon(coords)
    x, y = shapely.concave_hull(poly, ratio=ratio, allow_holes=holes).exterior.xy

    if small_palette_num:
        if palette == 'Cividis':
            palette = Cividis[small_palette_num]
        elif palette == 'Inferno':
            palette = Inferno[small_palette_num]
        elif palette == 'Magma':
            palette = Magma[small_palette_num]
        elif palette == 'Plasma':
            palette = Plasma[small_palette_num]
        elif palette == 'Viridis':
            palette = Viridis[small_palette_num]
        elif palette == 'Turbo':
            palette = Turbo[small_palette_num]
        elif palette == 'Grey':
            palette = Greys[small_palette_num]
    else:
        if palette == 'Cividis':
            palette = Cividis256
        elif palette == 'Inferno':
            palette = Inferno256
        elif palette == 'Magma':
            palette = Magma256
        elif palette == 'Plasma':
            palette = Plasma256
        elif palette == 'Viridis':
            palette = Viridis256
        elif palette == 'Turbo':
            palette = Turbo256
        elif palette == 'Grey':
            palette = Greys256

    # Define a color mapper based on the chosen palette and the range of z values
    if log_scale:
        mapper = log_cmap(field_name=z_column, palette=palette, low=min(df[z_column]), high=max(df[z_column]))
    else:
        mapper = linear_cmap(field_name=z_column, palette=palette, low=min(df[z_column]), high=max(df[z_column]))

    p = figure(
		x_axis_label=x_column,
		y_axis_label=y_column,
		aspect_ratio = 17.5 / 8,
		sizing_mode='scale_width',
		output_backend='webgl',
	)
    source = ColumnDataSource(df)
    p.circle(
        source=source,
        x=x_column,
        y=y_column,
        color=mapper
    )
    p.line(
        x=x,
        y=y,
        color='black'
    )

    if small_palette_num and not log_scale:
        n_ticks = np.linspace(df[z_column].min(), df[z_column].max(), small_palette_num + 1).round()
        color_bar = ColorBar(
            color_mapper=mapper['transform'],
            width=8,
            location=(0, 0),
            ticker=FixedTicker(ticks=n_ticks)
        )
    elif small_palette_num and log_scale:
        n_ticks = np.logspace(math.log(df[z_column].min()) / math.log(10), math.log(df[z_column].max()) / math.log(10), small_palette_num + 1).round()
        st.write(n_ticks)
        color_bar = ColorBar(
            color_mapper=mapper['transform'],
            width=8,
            location=(0, 0),
            ticker=FixedTicker(ticks=n_ticks)
        )
    else:
        color_bar = ColorBar(color_mapper=mapper['transform'], width=8, location=(0,0))

    p.add_layout(color_bar, 'right')

    t = f'''

        {x_column}:  @{x_column}

        {y_column}:  @{y_column}

        {z_column}:  @{z_column}

    '''
    num = 1
    f = CustomJSHover(
        code = f'''
        special_vars.indices = special_vars.indices.slice(0, {num})
        return special_vars.indices.includes(special_vars.index) ? " " : " hidden "
        ''')
    p.add_tools(HoverTool(tooltips = t, formatters = {f'@{x_column}': f}))

    p.toolbar.logo = None

    st.bokeh_chart(p)

    hull_df = pd.DataFrame({'Hull X Coordinates' : x, 'Hull Y Coordinates': y})
    st.download_button(
        label='Download Hull Coordinates',
        data=hull_df.to_csv(index=False).encode('utf-8'),
        file_name='Hull Coordinates.csv',
        mime='text/csv'
    )

def main() -> None:
    st.header('Convex Hull Graphing Tool')

    st.info('NOTE: This tool assumes your data is clean.')

    file = st.file_uploader(
        label='Upload your data',
        type=['csv', 'xlsx']
    )

    if file is None:
        st.info('Please upload some data to get started')
        return

    df = read_file(file)
    st.write(df.head())

    data_columns, additional_column, color_column = st.columns(3)
    with data_columns:
        x_column, y_column = get_plot_columns(df)
        sensitivity = st.number_input(
            label='Sensitivity level',
            min_value=0.0,
            max_value=1.0,
            value=0.02,
            format='%.4f'
        )
        holes = st.checkbox(
            label='Allow holes?',
            value=False
        )
    with additional_column:
        z_column = None
        if st.checkbox('Add a third column for color?'):
            z_column, palette = get_additional_column(df)

    with color_column:
        if z_column is not None:
            small_palette_num = None
            log_scale = st.checkbox('Use Log Scale for Z Column?')
            continuous_palette = st.checkbox('Use a Continuous palette? (Smooth color transition)', value=True)
            if not continuous_palette:
                small_palette_num = st.number_input(
                    label='How many unique colors?',
                    min_value=3,
                    value=8,
                    max_value=11
                )

    try:
        if x_column == y_column:
            st.warning('Your X and Y columns must be different')
            return

        if z_column is None:
            plot(df, x_column, y_column, sensitivity, holes)
        else:
            plot_contour(df, x_column, y_column, z_column, sensitivity, holes, palette, log_scale, small_palette_num)
    except:
        st.warning("We're sorry, something went wrong. Please make sure your data is clean")

if __name__ == '__main__':
    main()


"""
(c) Pavel Gurkov aka trueneu, 2016, All Rights Reserved

Adapter to mysql/business logic related to data schemes module (repacking queries/replies)
"""

import datetime
import hashlib
from w_is_for_wallpaper import mysqlpools as pools
from w_is_for_wallpaper import devices
from w_is_for_wallpaper import dict_repacker
from w_is_for_wallpaper import history
from w_is_for_wallpaper import entities
from w_is_for_wallpaper import mysql_db_themes_images
from w_is_for_wallpaper import responder

def get_coords_by_id_device_type(image_id, device_type):
    return pools.db_themes_images.get_image_coords_by_id_device_type(image_id, device_type)

def get_all_coords_by_id(image_id):
    return {'desktop': pools.db_themes_images.get_image_coords_by_id_device_type(image_id, devices.DeviceType.desktop),
            'cellphone': pools.db_themes_images.get_image_coords_by_id_device_type(image_id, devices.DeviceType.cellphone),
            'band': pools.db_themes_images.get_image_coords_by_id_device_type(image_id, devices.DeviceType.band)}

def image_reply_repack(image_info):
    return dict_repacker.repack({'id': 'id',
                                 'thumbnail': 'thumbnail', 'source': 'source', 'categories': 'categories',
                                 'license': 'license', 'colors_matched': 'colors_matched',
                                 'resolution': {'width': 'width', 'height': 'height'},
                                 'uid': 'uid', 'mtime': 'mtime',
                                 'origin': 'origin', 'votes': 'votes'},
                                image_info)

def image_query_repack(image_dict):
    return dict_repacker.repack({'thumbnail': 'thumbnail', 'source': 'source', 'categories': 'categories',
                                 'license': 'license', 'colors_matched': 'colors_matched',
                                 'width': 'resolution:width', 'height': 'resolution:height',
                                 'uid': 'uid', 'origin': 'origin', 'votes': 'votes'},
                                image_dict)

def category_query_repack(category_dict):
    return dict_repacker.repack({'name': 'name', 'image_id_cover': 'image_id_cover'}, category_dict)

def category_reply_repack(category_info):
    return dict_repacker.repack({'name': 'name', 'id': 'id', 'image_id_cover': 'image_id_cover'}, category_info)

def history_reply_repack(history_info):
    return dict_repacker.repack({'id': 'id', 'event_type': 'event_type', 'entity_type': 'entity_type',
                                 'entity_id': 'entity_id', 'date': 'date'}, history_info)

def get_responder_proxy(ids, entity_type):
    entity_to_db_themes_images_function_mapping = {
        entities.EntityTypes.image: pools.db_themes_images.get_images,
        entities.EntityTypes.license: pools.db_themes_images.get_licenses,
        entities.EntityTypes.category: pools.db_themes_images.get_categories
    }
    entity_to_repackers = {
        entities.EntityTypes.image: image_reply_repack,
        entities.EntityTypes.category: category_reply_repack
    }
    entity_to_amount = {
        entities.EntityTypes.image: 20000,
        entities.EntityTypes.license: None,
        entities.EntityTypes.category: None
    }
    amount = None
    if ids is not None and not isinstance(ids, list):  # if we were asked with a list, we'll return a list of dicts
        ids_list = [ids]
    else:
        if ids is None:
            amount = entity_to_amount[entity_type]
        ids_list = ids
    if amount is None:
        res = entity_to_db_themes_images_function_mapping[entity_type](ids_list)
    else:
        res = entity_to_db_themes_images_function_mapping[entity_type](ids_list, amount=amount)
    if not res:
        return []
    res_repacked = (list(map(entity_to_repackers[entity_type], res))
                    if entity_type in entity_to_repackers else res)

    if ids is not None and not isinstance(ids, list):  # if we were asked with a single id, we'll return a dict
        res_repacked = res_repacked[0]

    return res_repacked

# IMAGE FUNCTIONS

def get_images(image_ids):
    return get_responder_proxy(image_ids, entities.EntityTypes.image)

def get_random_image(orientation, category_id, color_id):
    try:
        image_id = pools.db_themes_images.get_random_image_id(orientation, category_id, color_id)['id']
    except (KeyError, AttributeError, TypeError):
        return None
    return get_responder_proxy(image_id, entities.EntityTypes.image)

def get_images_full_info_by_category_id(category_id):
    images_info = pools.db_themes_images.get_images_full_info_by_category(category_id)
    images_info_repacked = list(map(image_reply_repack, images_info))
    return images_info_repacked

def get_image_info_with_coords_by_id_device_type(image_id, device_type):
    image_info = pools.db_themes_images.get_image_info_by_id(image_id)
    image_info['coords'] = get_coords_by_id_device_type(image_id, device_type)
    return image_info

def get_newest_images_count_by_category(category_ids, since_date):
    if category_ids:
        return {'count': pools.db_themes_images.get_newest_images_count_by_category(category_ids, since_date)[
            'count(distinct image_id)']}
    else:
        return {'count': pools.db_themes_images.get_newest_images_count(since_date)[
            'count(distinct image_id)']}

# def get_images_full_info_newest(amount):
#     images_info = pools.db_themes_images.get_images_full_info_newest(amount)
#     images_info_repacked = list(map(image_reply_repack, images_info))
#     return images_info_repacked

def post_image(image_data):
    image_data_repacked = image_query_repack(image_data)
    category_ids = image_data_repacked.pop('categories')
    try:
        if len(get_categories(category_ids)) != len(category_ids):
            return None
    except TypeError:
        return None

    color_ids = image_data_repacked.pop('colors_matched')
    try:
        if len(get_colors(color_ids)) != len(color_ids):
            return None
    except TypeError:
        return None
    if not ('uid' in image_data_repacked):
        image_data_repacked['uid'] = hashlib.md5(image_data_repacked['source'].encode()).hexdigest()
    image_data_repacked['mtime'] = datetime.datetime.utcnow().strftime('%Y-%m-%d %H:%M:%S.%f')
    lastid = pools.db_themes_images.add_image(image_data_repacked)
    pools.db_themes_images.add_category_image_mapping(category_ids, lastid)
    pools.db_themes_images.add_color_image_mapping(color_ids, lastid)
    pools.db_themes_images.recalculate_image_top(lastid)
    # pools.db_themes_images.add_image_top_coefficient(lastid, None)  # TODO
    return {'id': lastid}

def get_themes_full_info():
    themes_info = pools.db_themes_images.get_themes_full_info()
    return themes_info

def delete_images(image_ids):
    if not isinstance(image_ids, list):
        image_ids = [image_ids]
    if len(get_images(image_ids)) != len(image_ids):
        return None

    pools.db_themes_images.delete_dailies_entries_by_images(image_ids)
    # categories_with_dailies = pools.db_themes_images.get_categories_images_are_daily_for(image_ids)
    # if categories_with_dailies:
    #     raise responder.ResponderException('Could not delete some images. Some are dailies. {}'.format(
    #         categories_with_dailies
    #     ))
    categories_with_covers = pools.db_themes_images.get_categories_images_are_cover_for(image_ids)
    if categories_with_covers:
        raise responder.ResponderException('Could not delete some images. Some are covers. {}'.format(
            categories_with_covers
        ))

    current_categories_ids = [x['category_id'] for x in get_categories_mapping(image_ids)]
    current_color_ids = [x['color_id'] for x in get_colors_mapping(image_ids)]
    pools.db_themes_images.delete_category_image_mapping(current_categories_ids, image_ids)
    pools.db_themes_images.delete_color_image_mapping(current_color_ids, image_ids)
    # pools.db_themes_images.delete_images_top_coefficient(image_ids)  # TODO
    res = pools.db_themes_images.delete_images(image_ids)

    return res

def get_categories_mapping(image_ids):
    return pools.db_themes_images.get_categories_mapping(image_ids)

def get_colors_mapping(image_ids):
    return pools.db_themes_images.get_colors_mapping(image_ids)

def modify_image(image_id, image_data):
    image_data_repacked = image_query_repack(image_data)

    try:
        category_ids = image_data_repacked.pop('categories', [])
        if category_ids and len(get_categories(category_ids)) != len(category_ids):
            return None

        color_ids = image_data_repacked.pop('colors_matched', [])
        if color_ids and len(get_colors(color_ids)) != len(color_ids):
            return None
    except TypeError:
        return None

    if not ('uid' in image_data_repacked) and 'source' in image_data_repacked:
        image_data_repacked['uid'] = hashlib.md5(image_data_repacked['source'].encode()).hexdigest()

    if not get_images(image_id):
        return None
    res = [None]

    # image_data_repacked['mtime'] = datetime.datetime.utcnow()  # TODO review ;we probably don't need it
                                                                # so we have rough info whether the image is old or not

    if image_data_repacked:
        res = pools.db_themes_images.modify_image(image_id, image_data_repacked)

    if category_ids:
        current_categories_ids_set = set([x['category_id'] for x in get_categories_mapping(image_id)])
        category_ids_set = set(category_ids)
        categories_ids_to_remove = list(current_categories_ids_set - category_ids_set)
        categories_ids_to_add = list(category_ids_set - current_categories_ids_set)
        pools.db_themes_images.delete_category_image_mapping(categories_ids_to_remove, image_id)
        pools.db_themes_images.add_category_image_mapping(categories_ids_to_add, image_id)

    if color_ids:
        current_color_ids_set = [x['color_id'] for x in get_colors_mapping(image_id)]
        pools.db_themes_images.delete_color_image_mapping(current_color_ids_set, image_id)
        pools.db_themes_images.add_color_image_mapping(color_ids, image_id)

    return res

def image_vote(image_id, votes_amount):
    res = pools.db_themes_images.modify_image_votes(image_id, votes_amount)
    return res

def get_image_top(amount):
    res = [x['id'] for x in pools.db_themes_images.get_image_top(amount)]
    return res

def recalculate_image_top():
    res = pools.db_themes_images.recalculate_images_top()
    return res

# CATEGORY FUNCTIONS

def get_categories(cat_id=None):
    return get_responder_proxy(cat_id, entities.EntityTypes.category)

def post_category(category_data):
    category_data_repacked = category_query_repack(category_data)
    return {'id': pools.db_themes_images.add_category(category_data_repacked)}

def delete_categories(category_ids):
    if not isinstance(category_ids, list):
        category_ids = [category_ids]
    if len(get_categories(category_ids)) != len(category_ids):
        return None
    pools.db_themes_images.delete_dailies_entries_by_categories(category_ids)
    return pools.db_themes_images.delete_categories(category_ids)

def modify_category(category_id, category_data):
    category_data_repacked = category_query_repack(category_data)
    if not get_categories(category_id):
        return None
    return pools.db_themes_images.modify_category(category_id, category_data_repacked)

def get_daily_image(category_id, daily_image_type, daily_image_orientation, daily_image_before_date):
    image_id = pools.db_themes_images.get_daily_image_id(category_id, daily_image_type, daily_image_orientation,
                                                         daily_image_before_date)
    if not image_id:
        return None
    image_info = get_images(image_id)
    return image_info

def delete_daily_image(category_id, daily_image_type, daily_image_orientation):
    res = pools.db_themes_images.delete_daily_image(category_id, daily_image_type, daily_image_orientation)
    return res

def modify_daily_image(category_id, daily_image_type, daily_image_orientation, daily_image_data):
    res = pools.db_themes_images.modify_daily_image(category_id, daily_image_type, daily_image_orientation, daily_image_data)
    return res

# COLOR FUNCTIONS

def get_colors(color_id=None):
    return pools.db_themes_images.get_colors(color_id)

def post_color(color_data):
    return {'id': pools.db_themes_images.add_color(color_data)}

def delete_colors(color_ids):
    if not isinstance(color_ids, list):
        color_ids = [color_ids]
    if len(get_colors(color_ids)) != len(color_ids):
        return None
    return pools.db_themes_images.delete_colors(color_ids)

def modify_color(color_id, color_data):
    if not get_colors(color_id):
        return None
    return pools.db_themes_images.modify_color(color_id, color_data)

# LICENSE FUNCTIONS

def get_licenses(license_ids):
    return get_responder_proxy(license_ids, entities.EntityTypes.license)

def post_license(license_data):
    return {'id': pools.db_themes_images.add_license(license_data)}

def delete_licenses(license_ids):
    if not isinstance(license_ids, list):
        license_ids = [license_ids]
    if len(get_licenses(license_ids)) != len(license_ids):
        return None
    return pools.db_themes_images.delete_licenses(license_ids)

def modify_license(license_id, license_data):
    if not get_licenses(license_id):
        return None
    return pools.db_themes_images.modify_license(license_id, license_data)

# HISTORY FUNCTIONS

def get_history(since_date, entity_type):
    entity_type_to_get_info_function = {
        entities.EntityTypes.category: get_categories,
        entities.EntityTypes.color: get_colors,
        entities.EntityTypes.license: get_licenses,
        entities.EntityTypes.image: get_images,
        entities.EntityTypes.theme: get_themes_full_info
    }

    history_data = pools.db_themes_images.get_history(since_date, entity_type)
    if entity_type == entities.EntityTypes.everything:  # raw history for everything
        res = list(map(history_reply_repack, history_data))
    else:
        try:
            last_event_date = history_data[-1]['date']
        except IndexError:
            last_event_date = None
        entity_ids_created = set([x['entity_id'] for x in history_data
                                  if x['event_type'] == history.EventTypes.create.value])
        entity_ids_modified = set([x['entity_id'] for x in history_data
                                  if x['event_type'] == history.EventTypes.modify.value])
        entity_ids_deleted = set([x['entity_id'] for x in history_data
                                  if x['event_type'] == history.EventTypes.delete.value])

        entity_ids_created_modified = list((entity_ids_created | entity_ids_modified) - entity_ids_deleted)

        entities_info_modified = entity_type_to_get_info_function[entity_type](
            entity_ids_created_modified
        )
        if entities_info_modified is None:
            entities_info_modified = []

        res = {'modified': entities_info_modified, 'deleted': list(entity_ids_deleted),
               'last_event_date': last_event_date}

    return res


import re

from django import forms
from django.contrib import admin
from django.core.exceptions import ValidationError

from cegs_portal.search.models import (
    Analysis,
    DNAFeature,
    Experiment,
    ExperimentDataFileInfo,
    Facet,
    FacetValue,
    File,
    Pipeline,
    RegulatoryEffectObservation,
)

def validate_ref_genome(value):
    if value != "" and not re.match(r"GRCh\d+(\.\d+)?", value):
        raise ValidationError(
            "%(ref)s is not a valid genome. Must be in the form of 'GRChXX' or 'GRChXX.Y' where Y is a patch level.",
            params={"ref": value},
        )

class DNAFeatureAdmin(admin.ModelAdmin):
    raw_id_fields = ("closest_gene", "parent", "experiment_accession")
    list_display = ("accession_id", "feature_type", "public", "archived")
    search_fields = ["experiment__accession_id"]

admin.site.register(DNAFeature, DNAFeatureAdmin)

class FileForm(forms.ModelForm):
    p_val_threshold = forms.FloatField(max_value=1.0, min_value=0.0, step_size=0.01, required=False)
    significance_measure = forms.CharField(max_length=2048, required=False)
    ref_genome = forms.CharField(max_length=10, validators=[validate_ref_genome], required=False)

    class Meta:
        model = File
        fields = (
            "filename",
            "description",
            "url",
            "experiment",
            "size",
            "category",
            "p_val_threshold",
            "significance_measure",
            "ref_genome",
        )
        widgets = {
            "description": forms.Textarea(attrs={"rows": 6, "columns": 90}),
            "significance_measure": forms.Textarea(attrs={"rows": 6, "columns": 90}),
        }

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        data_file_info = self.instance.data_file_info
        if data_file_info is not None:
            self.fields["p_val_threshold"].initial = data_file_info.p_value_threshold
            self.fields["significance_measure"].initial = data_file_info.significance_measure
            ref_genome = data_file_info.ref_genome
            ref_genome_patch = data_file_info.ref_genome_patch
            if ref_genome != "":
                if ref_genome_patch is not None and ref_genome_patch != "":
                    self.fields["ref_genome"].initial = f"{ref_genome}.{ref_genome_patch}"
                else:
                    self.fields["ref_genome"].initial = data_file_info.ref_genome

    def clean(self):
        p_val_threshold = self.cleaned_data["p_val_threshold"]
        significance_measure = self.cleaned_data["significance_measure"]
        ref_genome_full = self.cleaned_data["ref_genome"]
        if not (
            all(((p_val_threshold == "" or p_val_threshold is None), significance_measure == "", ref_genome_full == ""))
            or all((p_val_threshold != "", significance_measure != "", ref_genome_full != ""))
        ):
            raise ValidationError(
                "Must specify either all of (p val threshold, significance measure, ref genome) or none of them"
            )
        return super().clean()

    def save(self, commit=True):
        p_val_threshold = self.cleaned_data["p_val_threshold"]
        significance_measure = self.cleaned_data["significance_measure"]
        ref_genome_full = self.cleaned_data["ref_genome"]
        ref_genome_parts = ref_genome_full.split(".")
        ref_genome = ref_genome_parts[0]
        ref_genome_patch = "" if len(ref_genome_parts) == 1 else ref_genome_parts[1]

        if (p_val_threshold == "" or p_val_threshold is None) and significance_measure == "" and ref_genome_full == "":
            if self.instance.data_file_info is not None:
                self.instance.data_file_info.delete()
                self.instance.data_file_info = None
        else:
            if self.instance.data_file_info is not None:
                self.instance.data_file_info.p_value_threshold = p_val_threshold
                self.instance.data_file_info.significance_measure = significance_measure
                self.instance.data_file_info.ref_genome = ref_genome
                self.instance.data_file_info.ref_genome_patch = ref_genome_patch
                self.instance.data_file_info.save()
            else:
                data_file_info = ExperimentDataFileInfo(
                    p_value_threshold=p_val_threshold,
                    significance_measure=significance_measure,
                )
                data_file_info.ref_genome = ref_genome
                data_file_info.ref_genome_patch = ref_genome_patch
                data_file_info.save()
                self.instance.data_file_info = data_file_info
        return super().save(commit)

class FileAdmin(admin.ModelAdmin):
    form = FileForm
    list_display = ("truncate_filename", "experiment_name", "description")

    @admin.display(description="File Name")
    def truncate_filename(self, obj):
        filename = obj.filename
        if len(filename) <= 30:
            return filename

        return f"{filename[:15]}...{filename[-15:]}"

    @admin.display(description="Experiment")
    def experiment_name(self, obj):
        experiment = obj.experiment
        return obj.experiment.name if experiment is not None else "N/A"

admin.site.register(File, FileAdmin)

# StackedInline stackes the Experiment Data File fields vertically instead of horizontally(TabularInline).
# Extra  = 0 reduces repeat field sections
class FileInlineAdmin(admin.StackedInline):
    form = FileForm
    model = File
    extra = 0
    verbose_name = "File"

    @admin.display
    def other_file_info(self, obj):
        return f"{obj.filename}"

class ExperimentFacetValueInlineAdmin(admin.StackedInline):
    model = Experiment.facet_values.through
    extra = 0
    verbose_name = "Facet Value"

    @admin.display
    def facet_info(self, obj):
        return f"{obj.value} ({obj.facet.name})"

class ExperimentForm(forms.ModelForm):
    class Meta:
        widgets = {
            "description": forms.Textarea(attrs={"rows": 6, "columns": 90}),
        }

class ExperimentAdmin(admin.ModelAdmin):
    form = ExperimentForm
    inlines = [ExperimentFacetValueInlineAdmin, FileInlineAdmin]
    fields = [
        "public",
        "archived",
        "accession_id",
        "name",
        "description",
        "experiment_type",
        "biosamples",
    ]
    list_display = ("accession_id", "name", "public", "archived")
    list_filter = ["public", "archived"]
    search_fields = ["name", "description"]

admin.site.register(Experiment, ExperimentAdmin)

admin.site.register(Facet)

class PipelineForm(forms.ModelForm):
    class Meta:
        widgets = {
            "description": forms.Textarea(attrs={"rows": 6, "columns": 90}),
        }

class PipelineInlineForm(admin.StackedInline):
    model = Pipeline
    form = PipelineForm
    extra = 0

class AnalysisFacetValueInlineAdmin(admin.StackedInline):
    model = Analysis.facet_values.through
    extra = 0
    verbose_name = "Facet Value"

    @admin.display
    def facet_info(self, obj):
        return f"{obj.value} ({obj.facet.name})"

class AnalysisForm(forms.ModelForm):
    verbose_name_plural = "Analyses"

class AnalysisAdmin(admin.ModelAdmin):
    form = AnalysisForm
    inlines = [AnalysisFacetValueInlineAdmin, PipelineInlineForm]
    list_display = (
        "accession_id",
        "description",
        "experiment_info",
    )
    fields = [
        "public",
        "archived",
        "accession_id",
        "description",
        "when",
    ]
    search_fields = ["experiment__accession_id", "experiment__name"]

    # TODO: This is definitely a cause of n+1 queries
    @admin.display(description="Experiment")
    def experiment_info(self, obj):
        exper = obj.experiment

        return f"{exper.accession_id}: {exper.description[:15]}..."

    @admin.display(description="Description")
    def description(self, obj):
        description = obj.description
        if len(description) <= 30:
            return description

        return f"{description[:30]}..."

admin.site.register(Analysis, AnalysisAdmin)

class FacetValueAdmin(admin.ModelAdmin):
    search_fields = ["value", "facet__name"]
    list_display = ["facet_list_display"]

    @admin.display
    def facet_list_display(self, obj):
        return f"{obj.value} ({obj.facet.name})"

admin.site.register(FacetValue, FacetValueAdmin)

class RegEffectAdmin(admin.ModelAdmin):
    raw_id_fields = ("sources", "targets", "experiment_accession")
    list_display = ("accession_id", "experiment", "public", "archived")
    search_fields = ["experiment__accession_id"]

admin.site.register(RegulatoryEffectObservation, RegEffectAdmin)


"""
Author: LihuiChen
E-mail: lihuichen@126.com
Note: The metrics for reduced-rolution is the same with the matlat codes opened by [Vivone20].
      Metrics for full-resolution have a little different results from the codes opened by [Vivone20].

Refercence: PansharpeningToolver1.3 and Pansharpening Toolbox for Distribution

Pansharpening metrics: The same implementation of CC, SAM, ERGAS, Q2n as the one in Matlab codes publised by:
    [Vivone15]  G. Vivone, L. Alparone, J. Chanussot, M. Dalla Mura, A. Garzelli, G. Licciardi, R. Restaino, and L. Wald,
                "A Critical Comparison Among Pansharpening Algorithms", IEEE Transactions on Geoscience and Remote Sensing, vol. 53, no. 5, pp. 2565�2586, May 2015.
    [Vivone20]  G. Vivone, M. Dalla Mura, A. Garzelli, R. Restaino, G. Scarpa, M.O. Ulfarsson, L. Alparone, and J. Chanussot,
                "A New Benchmark Based on Recent Advances in Multispectral Pansharpening: Revisiting pansharpening with classical and
                emerging pansharpening methods",IEEE Geoscience and Remote Sensing Magazine, doi: 10.1109/MGRS.2020.3019315.
"""
from scipy.ndimage import sobel
import numpy as np
# from scipy import signal, ndimage, misc
import cv2
from numpy.linalg import norm
# from PIL import Image

EPS = 1e-7

##########################################################
# Full Reference metrics for Reduced Resolution Assesment
##########################################################

def SAM(ms,ps,degs = True):
    result = np.double(ps)
    target = np.double(ms)

    if result.shape != target.shape: raise ValueError('Result and target arrays must have the same shape!')

    bands = target.shape[2]
    rnorm = (result ** 2).sum(axis=2)
    tnorm = (target ** 2).sum(axis=2)

    dotprod = (result * target).sum(axis=2)
    rnorm = np.sqrt(rnorm*tnorm)
    tnorm = rnorm
    rnorm = np.maximum(rnorm, EPS)
    cosines = (dotprod / rnorm)
    sam2d = np.arccos(cosines)
    if degs:
        sam2d = np.rad2deg(sam2d)

    rnorm   = tnorm[tnorm>0]
    dotprod = dotprod[tnorm>0]
    sam     = np.arccos(dotprod/rnorm).mean()

    if degs:
        sam = np.rad2deg(sam)
    # sam2d[np.invert(np.isfinite(sam2d))] = 0.  # arccos(1.) -> NaN
    return sam

def CC(img1, img2):
    """SCC for 2D (H, W)or 3D (H, W, C) image; uint or float[0, 1]"""
    if not  img1.shape == img2.shape:
        raise ValueError('Input images must have the same dimensions.')
    img1_ = img1.astype(np.float64)
    img2_ = img2.astype(np.float64)
    if img1_.ndim == 2:
        return np.corrcoef(img1_.reshape(1, -1), img2_.rehshape(1, -1))[0, 1]
    elif img1_.ndim == 3:
        ccs = [np.corrcoef(img1_[..., i].reshape(1, -1), img2_[..., i].reshape(1, -1))[0, 1]
               for i in range(img1_.shape[2])]
        return np.mean(ccs)
    else:
        raise ValueError('Wrong input image dimensions.')

def sCC(ms, ps):
    ps_sobel = sobel(ps, mode='constant')
    ms_sobel = sobel(ms, mode='constant')
    return  (np.sum(ps_sobel*ms_sobel)/np.sqrt(np.sum(ps_sobel*ps_sobel))/np.sqrt(np.sum(ms_sobel*ms_sobel)))

def _qindex(img1, img2, block_size=8):
    """Q-index for 2D (one-band) image, shape (H, W); uint or float [0, 1]"""
    assert block_size > 1, 'block_size shold be greater than 1!'
    img1_ = img1.astype(np.float64)
    img2_ = img2.astype(np.float64)
    window = np.ones((block_size, block_size)) / (block_size**2)
    # window_size = block_size**2
    # filter, valid
    pad_topleft = int(np.floor(block_size/2))
    pad_bottomright = block_size - 1 - pad_topleft
    mu1 = cv2.filter2D(img1_, -1, window)[pad_topleft:-pad_bottomright, pad_topleft:-pad_bottomright]
    mu2 = cv2.filter2D(img2_, -1, window)[pad_topleft:-pad_bottomright, pad_topleft:-pad_bottomright]
    mu1_sq = mu1**2
    mu2_sq = mu2**2
    mu1_mu2 = mu1 * mu2

    sigma1_sq = cv2.filter2D(img1_**2, -1, window)[pad_topleft:-pad_bottomright, pad_topleft:-pad_bottomright] - mu1_sq
    sigma2_sq = cv2.filter2D(img2_**2, -1, window)[pad_topleft:-pad_bottomright, pad_topleft:-pad_bottomright] - mu2_sq
    sigma12 = cv2.filter2D(img1_ * img2_, -1, window)[pad_topleft:-pad_bottomright, pad_topleft:-pad_bottomright] - mu1_mu2

    # all = 1, include the case of simga == mu == 0
    qindex_map = np.ones(sigma12.shape)
    # sigma == 0 and mu != 0
    idx = ((sigma1_sq + sigma2_sq) == 0) * ((mu1_sq + mu2_sq) != 0)
    qindex_map[idx] = 2 * mu1_mu2[idx] / (mu1_sq + mu2_sq)[idx]
    # sigma !=0 and mu == 0
    idx = ((sigma1_sq + sigma2_sq) != 0) * ((mu1_sq + mu2_sq) == 0)
    qindex_map[idx] = 2 * sigma12[idx] / (sigma1_sq + sigma2_sq)[idx]
    # sigma != 0 and mu != 0
    idx = ((sigma1_sq + sigma2_sq)!=0) * ((mu1_sq + mu2_sq) != 0)
    qindex_map[idx] =((2 * mu1_mu2[idx]) * (2 * sigma12[idx])) / (
        (mu1_sq + mu2_sq)[idx] * (sigma1_sq + sigma2_sq)[idx])
    return np.mean(qindex_map)

def Q_AVE(img1, img2, block_size=8):
    """Q-index for 2D (H, W) or 3D (H, W, C) image; uint or float [0, 1]"""
    if not img1.shape == img2.shape:
        raise ValueError('Input images must have the same dimensions.')
    if img1.ndim == 2:
        return _qindex(img1, img2, block_size)
    elif img1.ndim == 3:
        qindexs = [_qindex(img1[..., i], img2[..., i], block_size) for i in range(img1.shape[2])]
        return np.array(qindexs).mean()
    else:
        raise ValueError('Wrong input image dimensions.')

def ERGAS(img_fake, img_real, scale=4):
    """ERGAS for 2D (H, W) or 3D (H, W, C) image; uint or float [0, 1].
    scale = spatial resolution of PAN / spatial resolution of MUL, default 4."""
    if not img_fake.shape == img_real.shape:
        raise ValueError('Input images must have the same dimensions.')
    img_fake_ = img_fake.astype(np.float64)
    img_real_ = img_real.astype(np.float64)
    if img_fake_.ndim == 2:
        mean_real = img_real_.mean()
        mean_real = mean_real**2
        mse = np.mean((img_fake_ - img_real_)**2)
        mean_real = np.maximum(mean_real, EPS)
        return 100.0 / scale * np.sqrt(mse / mean_real)
    elif img_fake_.ndim == 3:
        means_real = img_real_.reshape(-1, img_real_.shape[2]).mean(axis=0)
        means_real = means_real**2
        means_real = np.maximum(means_real, EPS)
        mses = ((img_fake_ - img_real_)**2).reshape(-1, img_fake_.shape[2]).mean(axis=0)
        return 100.0 / scale * np.sqrt((mses/means_real).mean())
    else:
        raise ValueError('Wrong input image dimensions.')

##########################################################
# Using Gaussian filter matched MTF to degrade HRMS images
##########################################################

if __name__ == '__main__':

# import numpy as np
# import os
# ArbRPN_dir = '/home/ser606/Documents/LihuiChen/ArbRPN_20200916/results/SR/RNN_RESIDUAL_BI_PAN_FB_MASK/QB-FIX-4/x4/'
# GT_dir = '/home/ser606/Documents/LihuiChen/PanSharp_dataset/QB/test/MTF/4bands/HRMS_npy'
# LRMS_dir = '/home/ser606/Documents/LihuiChen/PanSharp_dataset/QB/test/MTF/4bands/LRMS_npy'
# PAN_dir = '/home/ser606/Documents/LihuiChen/PanSharp_dataset/QB/test/MTF/4bands/LRPAN_npy'
# # save_root = '/home/ser606/Documents/LihuiChen/compare/extra/reduced_resolution/test'

# ArbRPN_files = os.listdir(ArbRPN_dir)
# ArbRPN_files.sort()
# GT_files = os.listdir(GT_dir)
# GT_files.sort()
# LRMS_files = os.listdir(LRMS_dir)
# LRMS_files.sort()
# PAN_files = os.listdir((PAN_dir))
# PAN_files.sort()
# cc = []
# sam = []
# ergas = []
# q_ave = []
# q2n = []
# for i in range(len(ArbRPN_files)):
#     ps = np.load(os.path.join(ArbRPN_dir, ArbRPN_files[i]))
#     ps = ps.astype(np.float)
#     gt = np.load(os.path.join(GT_dir, GT_files[i])).astype(np.float)
#     pan = np.load(os.path.join(PAN_dir, PAN_files[i])).astype(np.float)
#     # print('%s  ||  %s \n'%(ArbRPN_files[i], PAN_files[i]))
#     # print((ps.dtype))
#     cc.append(CC(ps, gt))
#     sam.append(SAM(gt, ps))
#     ergas.append(ERGAS(ps, gt))
#     q_ave.append(Q_AVE(ps, gt))
#     q2n.append(Q2n(gt, ps, 32, 32))

# print(mean(q2n))

    # import os
    # TFNET_dir = '/home/ser606/Documents/LihuiChen/ArbRPN_20200916/results/SR/PARABIRNN/QB-Vanilla-BiRNN-FR/x4'
    # LRMS_dir = '/home/ser606/Documents/LihuiChen/PanSharp_dataset/QB/test/MS_full_resolution/MS_npy'
    # PAN_dir = '/home/ser606/Documents/LihuiChen/PanSharp_dataset/QB/test/PAN_full_resolution/PAN_npy'
    # # save_root = '/home/ser606/Documents/LihuiChen/compare/extra/full_resolution/test'ArbRPN_files = os.listdir(ArbRPN_dir)
    # TFNET_files = os.listdir(TFNET_dir)
    # TFNET_files.sort()
    # # GT_files = os.listdir(GT_dir)
    # # GT_files.sort()
    # LRMS_files = os.listdir(LRMS_dir)
    # LRMS_files.sort()
    # PAN_files = os.listdir((PAN_dir))
    # PAN_files.sort()
    # Dl_results = []
    # Ds_results = []
    # Qnr_results = []
    # for i in range(len(TFNET_files)):
    #     print('processing the %d-th image.\n'%i)
    #     ps = np.load(os.path.join(TFNET_dir, TFNET_files[i]))
    #     ps = ps.astype(np.float)
    #     lrms = np.load(os.path.join(LRMS_dir, LRMS_files[i])).astype(np.float)
    #     pan = np.load(os.path.join(PAN_dir, PAN_files[i])).astype(np.float)
    #     # print('%s  ||  %s \n'%(ArbRPN_files[i], PAN_files[i]))
    #     # print((ps.dtype))
    #     msexp = upsample_mat_interp23(lrms, 4)
    #     dl, ds, hqnr = HQNR(ps, lrms, msexp, pan, 32, 'QB', 4)
    #     Dl_results.append(dl)
    #     Ds_results.append(ds)
    #     Qnr_results.append(hqnr)
    # print(sum(Qnr_results)/len(Qnr_results))
    a = np.random.randn(240, 240, 128)
    b = np.random.randn(240, 240, 128)
    Q2n(a,b)


from datetime import datetime
import numpy as np
import warnings
warnings.filterwarnings('ignore', category=FutureWarning)
import tensorflow as tf
import os
import cv2
import sys
import time
from tensorflow.contrib.framework.python.framework import checkpoint_utils
BASE_DIR = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
print("PID:", os.getpid())
print(os.path.join(BASE_DIR, 'models'))
sys.path.append(BASE_DIR)  # model
sys.path.append(os.path.join(BASE_DIR, 'models'))
sys.path.append(os.path.join(BASE_DIR, 'utils'))
sys.path.append(os.path.join(BASE_DIR, 'data_load'))
sys.path.append(os.path.join(BASE_DIR, 'preprocessing'))
import model_normalization as model
import data_ivt_h5_queue  # as data
import output_utils
import create_file_lst
sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), '../preprocessing'))
from normal_gen import save_norm
import gpu_create_manifold_ivt as ct
import argparse
import pymesh
import trimesh
import pandas as pd
import normal_gen
from sklearn.neighbors import DistanceMetric as dm
from sklearn.neighbors import NearestNeighbors
from random import sample
lst_dir, cats, all_cats, raw_dirs = create_file_lst.get_all_info()

slim = tf.contrib.slim
FLAGS=None
RESULT_PATH = None

LOG_FOUT = None

MAXOUT = None
NUM_INPUT_POINTS = None

BN_INIT_DECAY = 0.5
BN_DECAY_DECAY_RATE = 0.5
BN_DECAY_DECAY_STEP = None
BN_DECAY_CLIP = 0.99

def log_string(out_str):
    LOG_FOUT.write(out_str + '\n')
    LOG_FOUT.flush()
    print(out_str)

def get_bn_decay(batch):
    bn_momentum = tf.compat.v1.train.exponential_decay(
        BN_INIT_DECAY,
        batch * FLAGS.batch_size,
        BN_DECAY_DECAY_STEP,
        BN_DECAY_DECAY_RATE,
        staircase=True)
    bn_decay = tf.minimum(BN_DECAY_CLIP, 1 - bn_momentum)
    return bn_decay

class NoStdStreams(object):
    def __init__(self, stdout=None, stderr=None):
        self.devnull = open(os.devnull, 'w')
        self._stdout = stdout or self.devnull or sys.stdout
        self._stderr = stderr or self.devnull or sys.stderr

    def __enter__(self):
        self.old_stdout, self.old_stderr = sys.stdout, sys.stderr
        self.old_stdout.flush()
        self.old_stderr.flush()
        sys.stdout, sys.stderr = self._stdout, self._stderr

    def __exit__(self, exc_type, exc_value, traceback):
        self._stdout.flush()
        self._stderr.flush()
        sys.stdout = self.old_stdout
        sys.stderr = self.old_stderr
        self.devnull.close()

def load_model(sess, LOAD_MODEL_FILE, prefixs, strict=False):
    vars_in_pretrained_model = dict(checkpoint_utils.list_variables(LOAD_MODEL_FILE))
    # print(vars_in_pretrained_model)
    vars_in_defined_model = []

    for var in tf.trainable_variables():
        if isinstance(prefixs, list):
            for prefix in prefixs:
                if (var.op.name.startswith(prefix)) and (var.op.name in vars_in_pretrained_model.keys()) and ('logits' not in var.op.name):
                    if (list(var.shape) == vars_in_pretrained_model[var.op.name]):
                        vars_in_defined_model.append(var)
        else:
            if (var.op.name.startswith(prefixs)) and (var.op.name in vars_in_pretrained_model.keys()) and ('logits' not in var.op.name):
                if (list(var.shape) == vars_in_pretrained_model[var.op.name]):
                    vars_in_defined_model.append(var)
    # print(vars_in_defined_model)
    saver = tf.train.Saver(vars_in_defined_model)
    try:
        saver.restore(sess, LOAD_MODEL_FILE)
        print("Model loaded in file: %s" % (LOAD_MODEL_FILE))
    except:
        if strict:
            print("Fail to load modelfile: %s" % LOAD_MODEL_FILE)
            return False
        else:
            print("Fail loaded in file: %s" % (LOAD_MODEL_FILE))
            return True

    return True

def load_model_strict(sess, saver, restore_model):
    ckptstate = tf.train.get_checkpoint_state(restore_model)
    if ckptstate is not None:
        LOAD_MODEL_FILE = os.path.join(restore_model, os.path.basename(ckptstate.model_checkpoint_path))
        saver.restore(sess, LOAD_MODEL_FILE)
    return sess

def test(raw_dirs, TEST_DATASET):
    log_string(FLAGS.log_dir)
    with tf.Graph().as_default():
        with tf.device('/gpu:0'):
            input_pls = model.placeholder_inputs(scope='inputs_pl', FLAGS=FLAGS, num_pnts=NUM_INPUT_POINTS)
            is_training_pl = tf.compat.v1.placeholder(tf.bool, shape=())
            print(is_training_pl)
            batch = tf.Variable(0, name='batch')

            print("--- Get model and loss")
            # Get model and loss

            end_points = model.get_model(input_pls, is_training_pl, bn=False, FLAGS=FLAGS)
            loss, end_points = model.get_loss(end_points, FLAGS=FLAGS)
            gpu_options = tf.compat.v1.GPUOptions()  # (per_process_gpu_memory_fraction=0.99)
            config = tf.compat.v1.ConfigProto(gpu_options=gpu_options)
            config.gpu_options.allow_growth = True
            config.allow_soft_placement = True
            config.log_device_placement = False
            sess = tf.compat.v1.Session(config=config)

            ##### all
            # Init variables
            init = tf.compat.v1.global_variables_initializer()
            sess.run(init)

            ######### Loading Checkpoint ###############
            # Overall
            saver = tf.compat.v1.train.Saver([v for v in tf.compat.v1.get_collection_ref(tf.compat.v1.GraphKeys.GLOBAL_VARIABLES) if ('lr' not in v.name) and ('batch' not in v.name)])

            ###########################################

            ops = {'input_pls': input_pls,
                   'is_training_pl': is_training_pl,
                   'loss': loss,
                   'step': batch,
                   'saver': saver,
                   'end_points': end_points}
            sys.stdout.flush()

            TEST_DATASET.start()
            log_string('**** INFERENCE  ****')
            sys.stdout.flush()
            test_one_epoch(sess, ops, raw_dirs, TEST_DATASET)
            TEST_DATASET.shutdown()
            print("Done!")

def test_one_epoch(sess, ops, raw_dirs, TEST_DATASET):
    num_batches = int(len(TEST_DATASET) / FLAGS.batch_size)
    print('num_batches', num_batches)
    for batch_idx in range(num_batches):
        start_fetch_tic = time.time()
        batch_data = TEST_DATASET.fetch()
        view_id = batch_data["view_id"][0]
        outdir = os.path.join(FLAGS.outdir, batch_data["cat_id"][0], batch_data["obj_nm"][0], str(view_id), FLAGS.unitype)
        os.makedirs(outdir, exist_ok=True)
        nums = FLAGS.initnums
        weightform = FLAGS.weightform
        distr = FLAGS.distr

        sess = load_model_strict(sess,  ops["saver"], FLAGS.restore_model)

        if FLAGS.unitype == "uni":
            loc, norm, std, weights, gt_pc, tries, face_norms, vert_norms, dist = unisample_pnts(sess, ops, -1, batch_data, None, FLAGS.res, nums, threshold=FLAGS.uni_thresh, stdratio=FLAGS.stdratio, stdlwb=FLAGS.stdlwb[0], stdupb=FLAGS.stdupb[0])
            save_norm(loc, norm, os.path.join(outdir, "uni.ply"))
        elif FLAGS.unitype == "ball":
            loc, norm, std, weights, gt_pc, tries, face_norms, vert_norms, dist = ballsample_pnts(sess, ops, FLAGS.anglenums, -1, batch_data, None, nums, threshold=FLAGS.uni_thresh, stdratio=FLAGS.stdratio, stdlwb=FLAGS.stdlwb[0], stdupb=FLAGS.stdupb[0])
            save_norm(loc, norm, os.path.join(outdir, "ball.ply"))
        if FLAGS.restore_surfmodel != "":
            sess = load_model_strict(sess, ops["saver"], FLAGS.restore_surfmodel)

        for i in range(FLAGS.rounds):
            if FLAGS.rounds > 2 and i == (FLAGS.rounds - 2):
                weightform = "even"
                num_ratio = 1
            else:
                num_ratio = FLAGS.num_ratio
            nums = nums * num_ratio
            pc = sample_from_MM(loc, norm, std, weights, nums, num_ratio, dist, distr=distr, gridsize=FLAGS.gridsize)
            batch_data['pnts'] = np.array([pc])
            # print("batch_data['pnts'].shape", batch_data['pnts'].shape)
            loc, norm, std, weights, dist = nearsample_pnts(sess, ops, i, batch_data, tries, face_norms, vert_norms, stdratio=FLAGS.stdratio, weightform=weightform, stdlwb=FLAGS.stdlwb[i+1], stdupb=FLAGS.stdupb[i+1])
            np.savetxt(os.path.join(outdir, "surf_samp{}.txt".format(i)), pc, delimiter=";")
            save_norm(loc, norm, os.path.join(outdir, "surf_{}.ply".format(i)))
        sys.stdout.flush()
        print("finished inference {}/{}".format(batch_idx+1, num_batches))

def gt_test_one_epoch(raw_dirs, TEST_DATASET):
    num_batches = int(len(TEST_DATASET) / FLAGS.batch_size)
    print('num_batches', num_batches)
    log_string(FLAGS.log_dir)
    for batch_idx in range(num_batches):
        start_fetch_tic = time.time()
        batch_data = TEST_DATASET.fetch()
        view_id = batch_data["view_id"][0]
        outdir = os.path.join(FLAGS.outdir, batch_data["cat_id"][0], batch_data["obj_nm"][0], str(view_id), FLAGS.unitype)
        os.makedirs(outdir, exist_ok=True)
        nums = FLAGS.initnums
        weightform = FLAGS.weightform
        distr = FLAGS.distr
        modelfile = os.path.join(raw_dirs['norm_mesh_dir'], batch_data["cat_id"][0], batch_data["obj_nm"][0], "pc_norm.obj")

        if FLAGS.unitype == "uni":
            loc, norm, std, weights, gt_pc, tries, face_norms, vert_norms, dist = unisample_pnts(None, None, -1, batch_data, modelfile, FLAGS.res, nums, threshold=FLAGS.uni_thresh, stdratio=FLAGS.stdratio, stdlwb=FLAGS.stdlwb[0], stdupb=FLAGS.stdupb[0])
            save_norm(loc, norm, os.path.join(outdir, "uni.ply"))
        elif FLAGS.unitype == "ball":
            loc, norm, std, weights, gt_pc, tries, face_norms, vert_norms, dist = ballsample_pnts(None, None, -1,  FLAGS.anglenums, batch_data, modelfile, nums, threshold=FLAGS.uni_thresh, stdratio=FLAGS.stdratio, stdlwb=FLAGS.stdlwb[0], stdupb=FLAGS.stdupb[0])
            save_norm(loc, norm, os.path.join(outdir, "ball.ply"))

        for i in range(FLAGS.rounds):
            num_ratio = FLAGS.num_ratio
            nums = nums * num_ratio
            pc = sample_from_MM(loc, norm, std, weights, nums, num_ratio, dist, distr=distr, gridsize=FLAGS.gridsize)
            batch_data['pnts'] = np.array([pc])
            print("batch_data['pnts'].shape", batch_data['pnts'].shape)
            loc, norm, std, weights, dist = nearsample_pnts(None, None, i, batch_data, tries, face_norms, vert_norms, stdratio=FLAGS.stdratio, weightform=weightform, stdlwb=FLAGS.stdlwb[i+1], stdupb=FLAGS.stdupb[i+1])
            # print("pc.shape", pc.shape, "loc.shape", loc.shape, "std.shape", std.shape, "weights.shape", weights.shape)
            np.savetxt(os.path.join(FLAGS.outdir, "surf_samp{}.txt".format(i)), pc, delimiter=";")
            save_norm(loc, norm, os.path.join(outdir, "surf_{}.ply".format(i)))
        sys.stdout.flush()
        print("finished gt {}/{}".format(batch_idx+1, num_batches))
    print("Done!")

def rectify(norm, right_drct):
    cosim = np.dot(norm, right_drct)
    norm = -norm if cosim < 0 else norm
    return norm

def randomneg(norm):
    ind = np.array(sample([i for i in range(norm.shape[0])], norm.shape[0]//2))
    norm[ind] = -norm[ind]
    print(ind.shape, norm.shape)
    return norm

def get_right_drct(norm, ind, nnindices, nndistances, flag):
    marked_nn_ind = nnindices[flag[nnindices] > 0]
    if marked_nn_ind.shape[0] == 1:
        return norm[marked_nn_ind[0]]
    return np.mean(norm[marked_nn_ind], axis = 0)

def linearRect(loc, norm):
    indarr = np.arange(loc.shape[0])
    flag = np.zeros(loc.shape[0], dtype=int)
    dist = dm.get_metric("euclidean")
    dist_matrix = dist.pairwise(loc)
    nbrs = NearestNeighbors(n_neighbors=5, algorithm='ball_tree').fit(loc)
    nndistances, nnindices = nbrs.kneighbors(loc)
    # dist_sorted_ind = np.argsort(dist_matrix, axis=1)[:,1:]
    flag_closest_rected = np.ones_like(flag, dtype=int) * -1
    flag_closest_rected_dist = np.ones_like(flag, dtype=float) * 10
    ind = np.argmax(np.linalg.norm(loc, axis=1))
    right_drct = loc[ind] / np.linalg.norm(loc[ind])
    for i in range(loc.shape[0]):
        if i != 0:
            unmarkind = np.argmin(flag_closest_rected_dist[flag<1])
            ind = indarr[flag<1][unmarkind]
            right_drct = get_right_drct(norm, ind, nnindices[ind], nndistances[ind], flag)
            # print(ind, flag_closest_rected[ind], "flag_closest_rected_dist[ind]", flag_closest_rected_dist[ind])
        norm[ind] = rectify(norm[ind], right_drct)
        flag[ind] = 1
        flag_closest_rected_dist_new = np.minimum(dist_matrix[:, ind],flag_closest_rected_dist)
        flag_closest_rected = np.where(flag_closest_rected_dist_new<flag_closest_rected_dist, ind, flag_closest_rected)
        flag_closest_rected_dist = flag_closest_rected_dist_new
    return norm

def inference_batch(sess, ops, roundnum, batch_data):
    """ ops: dict mapping from string to tf ops """
    is_training = False
    log_string(str(datetime.now()))
    SPLIT_SIZE = int(np.ceil(batch_data['pnts'].shape[1] / NUM_INPUT_POINTS))
    totalnum = batch_data['pnts'].shape[1]
    extra_pts = np.zeros((1, SPLIT_SIZE * NUM_INPUT_POINTS - totalnum, 3), dtype=np.float32)
    batch_points = np.concatenate([batch_data['pnts'], extra_pts], axis = 1).reshape((SPLIT_SIZE, 1, -1, 3))
    pred_ivt_lst = []
    pred_drct_lst = []
    tic = time.time()
    for sp in range(SPLIT_SIZE):
        feed_dict = {ops['is_training_pl']: is_training,
                     ops['input_pls']['pnts']: batch_points[sp,:],
                     ops['input_pls']['imgs']: batch_data['imgs'],
                     ops['input_pls']['obj_rot_mats']: batch_data['obj_rot_mats'],
                     ops['input_pls']['trans_mats']: batch_data['trans_mats']}
        output_list = [ops['end_points']['pred_ivts_xyz'], ops['end_points']['pred_ivts_dist'],
                       ops['end_points']['pred_ivts_direction'], ops['end_points']['imgs']]
        loss_list = []
        outputs = sess.run(output_list + loss_list, feed_dict=feed_dict)
        pred_xyz_val, pred_dist_val, pred_direction_val, _ = outputs[:]
        pred_ivt_lst.append(pred_xyz_val)
        pred_drct_lst.append(pred_direction_val)
    pred_ivts = np.concatenate(pred_ivt_lst, axis=1).reshape(-1,3)[:totalnum,:]
    pred_directions = np.concatenate(pred_drct_lst, axis=1).reshape(-1,3)[:totalnum,:]
    outstr = ' -----rounds %d, %d points ------ ' % (roundnum, totalnum)
    outstr += ' time per obj: %.02f, ' % (time.time() - tic)
    log_string(outstr)
    return pred_ivts, pred_directions

def get_unigrid(ivt_res):
    grids = int(1 / ivt_res)
    x_ = np.linspace(-1, 1, num=grids).astype(np.float32)
    y_ = np.linspace(-1, 1, num=grids).astype(np.float32)
    z_ = np.linspace(-1, 1, num=grids).astype(np.float32)
    x, y, z = np.meshgrid(x_, y_, z_, indexing='ij')
    return np.stack([x.reshape(-1), y.reshape(-1), z.reshape(-1)], axis=1)

def get_normalized_mesh(model_file):
    total = 16384 * 5
    print("trimesh_load:", model_file)
    mesh_list = trimesh.load_mesh(model_file, process=False)
    if not isinstance(mesh_list, list):
        mesh_list = [mesh_list]
    area_sum = 0
    area_lst = []
    for idx, mesh in enumerate(mesh_list):
        area = np.sum(mesh.area_faces)
        area_lst.append(area)
        area_sum += area
    area_lst = np.asarray(area_lst)
    amount_lst = (area_lst * total / area_sum).astype(np.int32)
    points_all = np.zeros((0, 3), dtype=np.float32)
    all_face_normals = np.zeros((0, 3), dtype=np.float32)
    all_vert_normals = np.zeros((0, 3, 3), dtype=np.float32)
    all_tries = np.zeros((0, 3, 3), dtype=np.float32)
    for i in range(amount_lst.shape[0]):
        mesh = mesh_list[i]
        # print("start sample surface of ", mesh.faces.shape[0])
        points, index = trimesh.sample.sample_surface(mesh, amount_lst[i])
        vert_ind = mesh.faces.reshape(-1)
        # print("end sample surface")
        all_tries = np.concatenate([all_tries, mesh.vertices[vert_ind].reshape([-1, 3, 3])], axis=0)
        all_face_normals = np.concatenate([all_face_normals, mesh.face_normals], axis=0)
        all_vert_normals = np.concatenate([all_vert_normals, mesh.vertex_normals[vert_ind].reshape([-1, 3, 3])], axis=0)
        points_all = np.concatenate([points_all, points], axis=0)
    return points_all, all_tries, all_face_normals, all_vert_normals

def unisample_pnts(sess, ops, roundnum, batch_data, modelfile, res, num, threshold=0.1, stdratio=10, stdlwb=0.0, stdupb=0.1):
    gt_pnts, tries, face_norms, vert_norms = None, None, None, None
    unigrid = get_unigrid(res)
    if not FLAGS.unionly:
        inds = np.random.choice(unigrid.shape[0], num * 8)
        unigrid = unigrid[inds]  # uni_ivts = ct.gpu_calculate_ivt(unigrid, tries, gpu)
    batch_data["pnts"] = np.array([unigrid])
    if sess is None:
        gt_pnts, tries, face_norms, vert_norms = get_normalized_mesh(modelfile)
        uni_ivts, uni_inter_norm = ct.gpu_calculate_ivt_norm(batch_data["pnts"][0], tries, face_norms, vert_norms, int(FLAGS.gpu), tries.shape[0]>700)
        uni_dist = np.linalg.norm(uni_ivts, axis=1, keepdims=True)
        uni_drcts = uni_inter_norm if FLAGS.norminter else -uni_ivts / uni_dist
        uni_dist = uni_dist.reshape((-1))
        print("uni_ivts.shape, uni_drcts.shape", uni_ivts.shape, uni_drcts.shape)
    else:
        uni_ivts, uni_drcts = inference_batch(sess, ops, roundnum, batch_data)
        uni_drcts = -uni_drcts
        uni_dist = np.linalg.norm(uni_ivts, axis=1)
    ind = uni_dist <= threshold
    uni_drcts = uni_drcts[ind]
    dist, uni_place, std = cal_std_loc(unigrid[ind], uni_ivts[ind], stdratio, stdlwb=stdlwb, stdupb=stdupb)
    return uni_place, -uni_drcts, std, np.full(std.shape[0], 1), gt_pnts, tries, face_norms, vert_norms, dist

def get_ballgrid(angles_num):
    phi = np.linspace(0, np.pi, num=angles_num)
    theta = np.linspace(0, 2 * np.pi, num=angles_num)
    x = np.outer(np.sin(theta), np.cos(phi)).reshape(-1)
    y = np.outer(np.sin(theta), np.sin(phi)).reshape(-1)
    z = np.outer(np.cos(theta), np.ones_like(phi)).reshape(-1)
    return np.stack([x, y, z], axis=1)

def ballsample_pnts(sess, ops, roundnum, angles_num, batch_data, modelfile, num, threshold=0.1, stdratio=10, stdlwb=0.0, stdupb=0.1):
    gt_pnts, tries, face_norms, vert_norms = None, None, None, None
    ballgrid = get_ballgrid(angles_num)
    if not FLAGS.unionly:
        inds = np.random.choice(ballgrid.shape[0], num * 8)
        ballgrid = ballgrid[inds]
    batch_data["pnts"] = np.array([ballgrid])
    if sess is None:
        gt_pnts, tries, face_norms, vert_norms = get_normalized_mesh(modelfile)
        ball_ivts, ball_inter_norm = ct.gpu_calculate_ivt_norm(batch_data["pnts"][0], tries, face_norms, vert_norms, int(FLAGS.gpu), tries.shape[0]>700)
        ball_dist = np.linalg.norm(ball_ivts, axis=1, keepdims=True)
        ball_drcts = ball_inter_norm if FLAGS.norminter else -ball_ivts / ball_dist
        ball_dist = ball_dist.reshape((-1))
        print("ball_ivts.shape, ball_drcts.shape", ball_ivts.shape, ball_drcts.shape)
    else:
        ball_ivts, ball_drcts = inference_batch(sess, ops, roundnum, batch_data)
        ball_drcts=-ball_drcts
        ball_dist = np.linalg.norm(uni_ivts, axis=1)
    # ind = ball_dist <= threshold
    # ivts = ball_drcts[ind]
    dist, ball_place, std = cal_std_loc(ballgrid, ball_ivts, stdratio, stdlwb=stdlwb, stdupb=stdupb)
    return ball_place, ball_drcts, std, np.full(std.shape[0], 1), gt_pnts, tries, face_norms, vert_norms, dist

def cal_std_loc(pnts, ivts, stdratio, stdlwb=0.0, stdupb=0.1):
    loc = pnts + ivts
    # print("ivts.shape",ivts.shape)
    dist = np.linalg.norm(ivts, axis=1)
    print("stdupb,stdlwb",stdupb,stdlwb)
    std = np.minimum(stdupb, np.maximum(np.minimum(stdlwb, dist), dist / stdratio))
    # print("np.amax(std)",np.amax(std))
    return dist, loc, std

def sample_from_MM(locs, norm, std, weights, nums, num_ratio, dist, distr="gaussian", gridsize=-1.0):
    if np.sum(std) != 0:
        print(locs.shape, norm.shape, std.shape, weights.shape)
        times = nums*2 // locs.shape[0]
        if times > 0:
            locs = np.tile(locs, (times, 1))
            norm = np.tile(norm, (times, 1))
            std = np.tile(std, (times))
            dist = np.tile(dist, (times))
            weights = np.tile(weights, (times))
        print(locs.shape, norm.shape, std.shape, weights.shape)
        sampled_pnts = np.zeros((locs.shape[0], 3))
        if distr == "gaussian" or distr == "ball":
            std = np.tile(np.expand_dims(std, axis=1), (1, 3))
            for i in range(nums):
                ind = inds[i]
                if distr == "gaussian":
                    p_loc = np.random.normal(locs[ind], std[ind])
                elif distr == "ball":
                    p_loc = ball_sample(locs[ind], std[ind])
                sampled_pnts[i, :] = p_loc
        else:
            back = float(distr)
            sampled_pnts[:,:] = normback_sample(locs, norm, std, np.sqrt(np.maximum(0,2*np.multiply(dist,std)- np.square(std)))) # dist*0.8)
            print("max shift",np.amax(np.sqrt(2*np.multiply(dist,std)- np.square(std))))

        if gridsize > 0:
            print("weights.shape", weights.shape)
            weights = grid_weight(sampled_pnts, gridsize, weights)
            print("weights.shape", weights.shape)
        weights = weights / np.sum(weights)
        print("weights.shape", weights.shape)
        print("locs.shape", locs.shape)
        print("std.shape", std.shape)
        print("nums", nums)
        inds = np.random.choice(weights.shape[0], nums, p=weights)
        return sampled_pnts[inds]
    else:
        print("std == 0, no samples")
        return locs

def grid_weight(loc, gridsize, weights):
    minXYZ, maxXYZ = np.min(loc, axis=0), np.max(loc, axis=0)
    xyzDims = np.floor((maxXYZ - minXYZ) / gridsize)
    norm_loc = loc - minXYZ
    norm_loc_xyzind = np.floor(norm_loc / gridsize)
    norm_loc_ind = norm_loc_xyzind[:,0] * xyzDims[1] * xyzDims[2] + norm_loc_xyzind[:,1] * xyzDims[2] + norm_loc_xyzind[:,2]
    unique, reverse_index, counts = np.unique(norm_loc_ind, return_inverse=True, return_counts=True)
    weight_count = counts[reverse_index]
    weights = weights / (1.0 * weight_count)
    return weights

def ball_sample(xyzmean, radius):
    uvw = np.random.normal(np.array([0, 0, 0]), np.array([1, 1, 1]))
    e = np.random.exponential(0.5)
    denom = (e + np.dot(uvw, uvw)) ** 0.5
    return xyzmean + uvw / denom * radius

def normback_sample(xyzmean, norm, back, radius):
    xyzmean = xyzmean + np.multiply(norm,np.expand_dims(back, axis=1))
    num = xyzmean.shape[0]
    r = np.multiply(radius, np.sqrt(np.random.uniform(0, 1, size=num)))
    a = np.random.uniform(0, np.pi * 2, size=num)
    xyzround = np.expand_dims(np.stack([np.multiply(r, np.cos(a)), np.multiply(r, np.sin(a)), np.zeros_like(r)], axis=1), axis=2)
    R = z_norm_matrix(norm)
    print(R.shape, xyzround.shape)
    xyzround = np.matmul(R, xyzround).reshape((-1,3))
    return xyzmean + xyzround

def z_norm_matrix(norm):
    bs = norm.shape[0]
    z_b = np.repeat(np.array([(0., 0., 1.)]), bs, axis=0)
    normal = np.cross(z_b, norm)
    print("normal of norm rotate:", normal.shape)
    sinal = np.linalg.norm(normal,axis=1,keepdims=True)
    cosal = np.sum(np.multiply(z_b, norm),axis=1,keepdims=True)
    # sinalsqrrev = 1/np.square(np.maximum(1e-5, sinal))
    ux = normal[:,[0]]
    uy = normal[:,[1]]
    uz = normal[:,[2]]
    zr = np.zeros_like(ux)
    W = np.concatenate([zr, -uz, uy, uz, zr, -ux, -uy, ux, zr], axis=1).reshape((-1,3,3))
    I = np.repeat(np.expand_dims(np.identity(3), axis = 0), bs, axis=0)
    Wsqr = np.matmul(W, W)
    C = 1/np.maximum(1e-5,1+cosal)
    R = I + W + np.expand_dims(C, axis=2) * Wsqr
    return R

def nearsample_pnts(sess, ops, roundnum, batch_data, tries, face_norms, vert_norms, stdratio=10, weightform="reverse", stdlwb=0.04, stdupb=0.1):
    if sess is None:
        surface_ivts, surf_inter_norm = ct.gpu_calculate_ivt_norm(batch_data["pnts"][0], tries, face_norms, vert_norms, int(FLAGS.gpu), tries.shape[0]>700)
        dist = np.linalg.norm(surface_ivts, axis=1, keepdims=True)
        surf_norm = surf_inter_norm if FLAGS.norminter else -surface_ivts / dist
        print("surface_ivts.shape, surface_ivts.shape", surface_ivts.shape, surface_ivts.shape)
    else:
        surface_ivts, surf_norm = inference_batch(sess, ops, roundnum, batch_data)
        surf_norm=-surf_norm
    # print("batch_data[pnts][0], surface_ivts",batch_data["pnts"][0].shape, surface_ivts.shape)
    dist, surface_place, std = cal_std_loc(batch_data["pnts"][0], surface_ivts, stdratio, stdlwb=stdlwb, stdupb=stdupb)
    if weightform == "even":
        weights = np.full(std.shape[0], 1)
    elif weightform == "reverse":
        weights = 1.0 / np.maximum(dist, 5e-3)

    return surface_place, surf_norm, std, weights, dist

if __name__ == "__main__":
    # nohup python -u inference.py --gpu 3 --img_feat_onestream --category chair  --restore_model ../train/checkpoint/onestream_small_spheregrid/chair_vgg_16_010000 --restore_surfmodel ../train/checkpoint/onestream_small_surf/chair_vgg_16_010000 --outdir  chair_drct_even_surfonly_uni --unionly --unitype uni --XYZ &> global_direct_chair_surf_evenweight_uni.log &

    # nohup python -u inference.py --gt --outdir  gt_noerrBall --unionly --stdupb 0.3 0.3 0.2 0.1 0.05 --stdlwb 0.01 0.01 0.01 0.01 0.01 &> gt_uni.log &

    #   nohup python -u inference.py --gt --outdir  gt_shirt1 --unionly --norminter --stdupb 0.3 0.3 0.2 0.1 0.05 --stdlwb 0.01 0.0 0.0 0.0 0.0 &> shirt.log &

    # full set

    parser = argparse.ArgumentParser()
    parser.add_argument('--gpu', type=str, default='0', help='GPU to use [default: GPU 0]')
    parser.add_argument('--category', type=str, default="all", help='Which single class to train on [default: None]')
    parser.add_argument('--log_dir', default='checkpoint', help='Log dir [default: log]')
    parser.add_argument('--num_pnts', type=int, default=0, help='Point Number [default: 2048]')
    parser.add_argument('--uni_num', type=int, default=0, help='Point Number [default: 2048]')
    parser.add_argument('--sphere_num', type=int, default=0, help='Point Number [default: 2048]')
    parser.add_argument('--num_classes', type=int, default=1024, help='vgg dim')
    parser.add_argument("--beta1", type=float, dest="beta1", default=0.5, help="beta1 of adams")
    parser.add_argument('--batch_size', type=int, default=1, help='Batch Size during training [default: 32]')
    parser.add_argument('--img_h', type=int, default=137, help='Image Height')
    parser.add_argument('--img_w', type=int, default=137, help='Image Width')
    parser.add_argument('--learning_rate', type=float, default=1e-5, help='Initial learning rate [default: 0.001]')
    parser.add_argument('--momentum', type=float, default=0.9, help='Initial learning rate [default: 0.9]')
    parser.add_argument('--optimizer', default='adam', help='adam or momentum [default: adam]')
    parser.add_argument('--restore_model', default='', help='restore_model')  # checkpoint/sdf_2d3d_sdfbasic2_nowd
    parser.add_argument('--restore_surfmodel', default='',
                        help='restore_model surface only')  # ../models/CNN/pretrained_model/vgg_16.ckpt
    parser.add_argument('--train_lst_dir', default=lst_dir, help='train mesh data list')
    parser.add_argument('--valid_lst_dir', default=lst_dir, help='test mesh data list')
    parser.add_argument('--decay_step', type=int, default=200000, help='Decay step for lr decay [default: 200000]')
    parser.add_argument('--decay_rate', type=float, default=0.9, help='Decay rate for lr decay [default: 0.7]')
    parser.add_argument('--weight_type', type=str, default="ntanh")
    parser.add_argument('--dir_loss_weight', type=float, default=0.1)
    parser.add_argument('--img_feat_onestream', action='store_true')
    parser.add_argument('--img_feat_twostream', action='store_true')
    parser.add_argument('--binary', action='store_true')
    parser.add_argument('--alpha', action='store_true')
    parser.add_argument('--augcolorfore', action='store_true')
    parser.add_argument('--augcolorback', action='store_true')
    parser.add_argument('--backcolorwhite', action='store_true')
    parser.add_argument('--rot', action='store_true')
    parser.add_argument('--XYZ', action='store_true')
    parser.add_argument('--cam_est', action='store_true')
    parser.add_argument('--cat_limit', type=int, default=168000, help="balance each category, 1500 * 24 = 36000")
    parser.add_argument('--multi_view', action='store_true')
    parser.add_argument('--test_lst_dir', default=lst_dir, help='test mesh data list')
    parser.add_argument('--rounds', type=int, default=4, help='how many rounds of ivf')
    parser.add_argument('--uni_thresh', type=float, default=0.1, help='threshold for uniform sampling')
    parser.add_argument('--res', type=float, default=0.01, help='cube resolution')
    parser.add_argument('--anglenums', type=int, default=200, help='angle resolution')
    parser.add_argument('--max_epoch', type=int, default=1, help='angle resolution')
    parser.add_argument('--initnums', type=int, default=8096, help='initial sampled uni point numbers')
    parser.add_argument('--num_ratio', type=int, default=2, help='point numbers expansion each round')
    parser.add_argument('--stdratio', type=int, default=4, help='')
    parser.add_argument('--stdupb', nargs='+', action='store', default=[0.1, 0.1, 0.03, 0.01, 0])
    parser.add_argument('--stdlwb', nargs='+', action='store', default=[0.08, 0.04, 0.003, 0.001, 0])
    # parser.add_argument('--stdupb', nargs='+', action='append', default=[0, 0])
    # parser.add_argument('--stdlwb', nargs='+', action='append', default=[0, 0])
    parser.add_argument('--distr', type=str, default="0.1", help='local points sampling: gaussian or ball')
    parser.add_argument('--weightform', type=str, default="even", help='GMM weight: even or reverse')
    parser.add_argument('--gridsize', type=float, default=0.02, help='GMM weight: even or reverse')
    parser.add_argument('--outdir', type=str, default="./", help='out_dir')
    parser.add_argument('--unitype', type=str, default="ball", help='out_dir')
    parser.add_argument('--unionly', action='store_true')
    parser.add_argument('--gt', action='store_true')
    parser.add_argument('--norminter', action='store_true')
    parser.add_argument('--smaller', action='store_true')
    parser.add_argument('--act', type=str, default="relu")
    parser.add_argument('--encoder', type=str, default='vgg_16',
                        help='encoder model: vgg_16, resnet_v1_50, resnet_v1_101, resnet_v2_50, resnet_v2_101')
    parser.add_argument('--wd', type=float, default=1e-6, help='Initial learning rate [default: 0.001]')
    parser.add_argument('--decoderskip', action='store_true')
    parser.add_argument('--lossw', nargs='+', action='store', default=[0.0, 1.0, 0.0, 0.0, 1.0, 0.0],
                        help="xyz, locnorm, locsqrnorm, dist, dirct, drct_abs")
    parser.add_argument('--distlimit', nargs='+', action='store', type=str,
                        default=[1.0, 0.9, 0.9, 0.8, 0.8, 0.7, 0.7, 0.6, 0.6, 0.5, 0.5, 0.4, 0.4, 0.3, 0.3, 0.2, 0.2,
                                 0.18, 0.18, 0.16, 0.16, 0.14, 0.14, 0.12, 0.12, 0.1, 0.1, 0.08, 0.08, 0.06, 0.06, 0.05,
                                 0.05, 0.04, 0.04, 0.03, 0.03, 0.02, 0.02, 0.01, 0.01, -0.01])
    parser.add_argument('--surfrange', nargs='+', action='store', default=[0.0, 0.15], help="lower bound, upperbound")
    FLAGS = parser.parse_args()
    FLAGS.stdupb = [float(i) for i in FLAGS.stdupb]
    FLAGS.stdlwb = [float(i) for i in FLAGS.stdlwb]
    FLAGS.lossw = [float(i) for i in FLAGS.lossw]
    FLAGS.distlimit = [float(i) for i in FLAGS.distlimit]
    FLAGS.surfrange = [float(i) for i in FLAGS.surfrange]
    print(FLAGS)

    # testarray = normback_sample(np.zeros((100,3)), np.tile(np.array([[0, 0.5, np.sqrt(3)*0.5]]), (100,1)), np.ones(100), np.ones(100))
    #
    # print(testarray)
    # np.savetxt("./test.txt", testarray, delimiter=";")
    # a = np.array([[[1,1,1],[2,2,2],[3,3,3]],[[5,5,5],[6,6,6],[7,7,7]],[[9,9,9],[11,11,11],[13,13,13]], [[93,90,90],[110,110,110],[130,130,130]]])
    # b = np.array([1,3])
    # print(np.take(a, b, axis=0))
    #
    # exit()

    os.environ["CUDA_VISIBLE_DEVICES"] = FLAGS.gpu

    if not os.path.exists(FLAGS.log_dir): os.makedirs(FLAGS.log_dir)

    RESULT_PATH = os.path.join(FLAGS.log_dir, 'test_results')
    if not os.path.exists(RESULT_PATH): os.mkdir(RESULT_PATH)

    os.system('cp %s.py %s' % (os.path.splitext(model.__file__)[0], FLAGS.log_dir))
    os.system('cp inference.py %s' % (FLAGS.log_dir))
    LOG_FOUT = open(os.path.join(FLAGS.log_dir, 'log_test.txt'), 'w')
    LOG_FOUT.write(str(FLAGS) + '\n')

    MAXOUT = FLAGS.initnums * FLAGS.num_ratio ** (FLAGS.rounds - 1)
    NUM_INPUT_POINTS = min(MAXOUT, 274625)
    print("NUM_INPUT_POINTS", NUM_INPUT_POINTS)

    BN_DECAY_DECAY_STEP = float(FLAGS.decay_step)

    if FLAGS.category != "all":
        cats = {
            FLAGS.category:cats[FLAGS.category]
        }

    TEST_LISTINFO = []
    cats_limit = {}

    cat_ids = []
    if FLAGS.category == "all":
        for key, value in cats.items():
            cat_ids.append(value)
            cats_limit[value] = 0
    else:
        cat_ids.append(cats[FLAGS.category])
        cats_limit[cats[FLAGS.category]] = 0

    for cat_id in cat_ids:
        test_lst = os.path.join(FLAGS.test_lst_dir, cat_id + "_test.lst")
        with open(test_lst, 'r') as f:
            lines = f.read().splitlines()
            for line in lines:
                for render in range(24):
                    cats_limit[cat_id] += 1
                    TEST_LISTINFO += [(cat_id, line.strip(), render)]

    # TEST_LISTINFO += [("03001627", "17e916fc863540ee3def89b32cef8e45", 11)]
    # TEST_LISTINFO += [("03001627", "1be38f2624022098f71e06115e9c3b3e", 0)]
    # TEST_LISTINFO += [("02691156", "1beb0776148870d4c511571426f8b16d", 0)]

    info = {'rendered_dir': raw_dirs["renderedh5_dir"],
            'ivt_dir': raw_dirs["ivt_mani_dir"]}
    if FLAGS.cam_est:
        info['rendered_dir'] = raw_dirs["renderedh5_dir_est"]

    print(info)

    TEST_DATASET = data_ivt_h5_queue.Pt_sdf_img(FLAGS, listinfo=TEST_LISTINFO, info=info, cats_limit=cats_limit, shuffle=False)

    # batch_data={"cat_id":["02691156"],"obj_nm":["1beb0776148870d4c511571426f8b16d"]}
    # model_file = os.path.join("/hdd_extra1/datasets/ShapeNet/march_cube_objs_v1/", batch_data["cat_id"][0], batch_data["obj_nm"][0], "isosurf.obj")
    #
    # command_str = "cp " + model_file + " ./" + FLAGS.outdir + "/"
    # print("command:", command_str)
    # os.system(command_str)
    # batch_data = {"cat_id": ["000"], "obj_nm": ["shirt3"]}
    # model_file = "./shirt.obj"
    # ct.get_normalize_mesh(model_file, "./", "./", "./", 0)
    # model_file = "./pc_norm.obj"
    # batch_data["model_file"] = model_file
    os.makedirs(FLAGS.outdir, exist_ok=True)

    if FLAGS.gt:
        gt_test_one_epoch(raw_dirs, TEST_DATASET)
    else:
        test(raw_dirs, TEST_DATASET)
    print("done!")
    LOG_FOUT.close()


import sqlite3

def create_tables():
    conn = sqlite3.connect("database.db")
    cursor = conn.cursor()

    # Verifica se a tabela 'clients' já existe
    cursor.execute("SELECT name FROM sqlite_master WHERE type='table' AND name='clients'")
    clients_table_exists = cursor.fetchone()

    if not clients_table_exists:
        # Criação da tabela 'clients'
        cursor.execute(
            """
            CREATE TABLE clients (
                id_client INTEGER PRIMARY KEY AUTOINCREMENT,
                name TEXT NOT NULL,
                cpf TEXT NOT NULL,
                telephone TEXT NOT NULL
            )
            """
        )

    # Verifica se a tabela 'motorcycles' já existe
    cursor.execute("SELECT name FROM sqlite_master WHERE type='table' AND name='motorcycles'")
    motorcycles_table_exists = cursor.fetchone()

    if not motorcycles_table_exists:
        # Criação da tabela 'motorcycles'
        cursor.execute(
            """
            CREATE TABLE motorcycles (
                id_motorcycle INTEGER PRIMARY KEY AUTOINCREMENT,
                brand TEXT NOT NULL,
                model TEXT NOT NULL,
                year INTEGER NOT NULL,
                price DECIMAL(10,2) NOT NULL
            )
            """
        )

        # Verifica se a tabela 'sales' já existe
        cursor.execute("SELECT name FROM sqlite_master WHERE type='table' AND name='sales'")
        sales_table_exists = cursor.fetchone()

        if not sales_table_exists:
            # Criação da tabela 'sales'
            cursor.execute(
                """
                CREATE TABLE sales (
                    id_sale INTEGER PRIMARY KEY AUTOINCREMENT,
                    id_client INTEGER NOT NULL,
                    id_motorcycle INTEGER NOT NULL,
                    sale_date TEXT NOT NULL,
                    FOREIGN KEY (id_client) REFERENCES clients (id_client),
                    FOREIGN KEY (id_motorcycle) REFERENCES motorcycles (id_motorcycle)
                )
                """
            )

    conn.commit()
    conn.close()

def get_db():
    conn = sqlite3.connect("database.db")
    return conn


# @Time    : 2022/7/14 下午5:48
# @Author  : Boyang
# @Site    :
# @File    : augmentation.py
# @Software: PyCharm
from __future__ import annotations

import os
from abc import ABC, abstractmethod
from concurrent.futures.process import ProcessPoolExecutor
from os.path import join as path_join
from typing import List, Optional, Callable

import albumentations as A
import cv2
import numpy as np
from albumentations.core.bbox_utils import BboxProcessor

from .common import chunk_dataset
from ..core.convertor import Convertor
from ..core.dataset import Dataset, ImageLabel, Box
from ..core.rotate import rotated_calculate_divide_point_bbox, rotate_image
from ..utils.imgutil import ImageUtil

MAX_WORKER = 5

class Invoker:
    def __init__(self):
        self.commands: List[Command] = []

    def add_command(self, c: Command):
        self.commands.append(c)

    def execute(self, ds: Dataset):
        for img_data, labels in chunk_dataset(ds):
            for c in self.commands:
                img_data, labels = c.execute(img_data, labels)

class Command(ABC):
    """Abstract Command"""

    @abstractmethod
    def execute(self, img_data: List[np.ndarray], labels: List[ImageLabel]):
        raise NotImplementedError

class Rotate(Command):
    """Concrete Command"""

    def __init__(self, receiver: Augmentation, degree: int, divide_num: int):
        self.receiver = receiver
        self.degree = degree
        self.ndivd = divide_num

    def execute(self, img_data: List[np.ndarray], labels: List[ImageLabel]):
        img_data, labels = self.receiver.rotate(self.degree, self.ndivd, img_data, labels)
        return img_data, labels

class PixelAugmentCommand(Command):

    def __init__(self, aug_func: Callable, min_area: float, min_visibility: float):
        self.aug_func = aug_func
        self.min_area = min_area
        self.min_visibility = min_visibility

    def execute(self, img_data: List[np.ndarray], labels: List[ImageLabel]):
        res_img = []
        for img in img_data:
            data = self.aug_func(image=img)
            res_img.append(data["image"])

        return res_img, labels

class SpatialAugmentCommand(Command):

    def __init__(self, aug_func: Callable, min_area: float, min_visibility: float):
        self.aug_func = aug_func
        self.min_area = min_area
        self.min_visibility = min_visibility

    def execute(self, img_data: List[np.ndarray], labels: List[ImageLabel]):
        process = BboxProcessor(
            A.BboxParams('pascal_voc', min_area=self.min_area, min_visibility=self.min_visibility))
        res_lbl, res_img = [], []
        for img, lbl in zip(img_data, labels):
            data = {"image": img, "bboxes": [[*b.coord, b.name] for b in lbl.boxes]}
            process.preprocess(data)
            data = self.aug_func(**data)
            process.postprocess(data)

            res_img.append(data["image"])
            boxes = [Box(*b) for b in data["bboxes"]]
            lbl.boxes = boxes
            res_lbl.append(lbl)

        return res_img, res_lbl

class Blur(PixelAugmentCommand):
    def __init__(self, blur_limit: int = 7, min_area: float = 0.0, min_visibility: float = 0.0, p: float = 1.0):
        aug_func = A.Blur(blur_limit, p=p)
        super(Blur, self).__init__(aug_func, min_area=min_area, min_visibility=min_visibility)

class CenterCrop(SpatialAugmentCommand):
    def __init__(self, width: int, height: int, min_area=0.0, min_visibility=0.0, p: float = 1.0):
        center_crop = A.CenterCrop(height, width, p=p)
        super(CenterCrop, self).__init__(center_crop, min_area, min_visibility)

class CoarseDropout(PixelAugmentCommand):
    def __init__(self,
                 min_area: float = 0.0,
                 min_visibility: float = 0.0,
                 max_holes: int = 8,
                 max_height: int = 8,
                 max_width: int = 8,
                 min_holes: Optional[int] = None,
                 min_height: Optional[int] = None,
                 min_width: Optional[int] = None,
                 fill_value: int = 0,
                 mask_fill_value: Optional[int] = None,
                 always_apply: bool = False,
                 p: float = 0.5
                 ):
        aug_func = A.CoarseDropout(max_holes, max_height, max_width, min_holes, min_height, min_width,
                                   fill_value, mask_fill_value, always_apply, p=p)

        super(CoarseDropout, self).__init__(aug_func, min_area, min_visibility)

class DownScale(PixelAugmentCommand):
    def __init__(self,
                 min_area: float = 0.0,
                 min_visibility: float = 0.0,
                 scale_min=0.25,
                 scale_max=0.25,
                 interpolation=cv2.INTER_NEAREST,
                 p: float = 0.5
                 ):
        aug_func = A.Downscale(scale_min=scale_min, scale_max=scale_max, interpolation=interpolation, p=p)
        super(DownScale, self).__init__(aug_func, min_area, min_visibility)

# class TestCommand(Command):
#     def __init__(self):
#         center_crop = A.CenterCrop(400, 400)
#         coarse_dropout = A.CoarseDropout()
#
#         self.compose = A.Compose([center_crop, coarse_dropout],
#                                  bbox_params=A.BboxParams('pascal_voc'))
#
#     def execute(self, img_data: List[np.ndarray], labels: List[ImageLabel]):
#         for img, lbl in zip(img_data, labels):
#             bboxes = [[*b.coord, b.name] for b in lbl.boxes]
#             self.compose(image=img, bboxes=bboxes)

class Save(Command):
    def __init__(self, receiver: Augmentation):
        self.receiver = receiver

    def execute(self, img_data: List[np.ndarray], labels: List[ImageLabel]):
        self.receiver.save(img_data, labels)
        return img_data, labels

class Augmentation:
    """Receiver"""

    def __init__(self, output: str, dst_dtype: str, contiguous: bool):
        self.output = output
        assert not os.path.exists(self.output), f"Directory `{self.output}` exist."

        self.convertor = Convertor.from_type(dst_dtype)
        self.convertor.create_dataset(*os.path.split(output))

        self.contiguous = contiguous

        self._count = 0
        self.pro_executor = ProcessPoolExecutor(max_workers=MAX_WORKER)

    @classmethod
    def rotate(cls, degree: int, number_divide: int, img_data: List[np.ndarray], labels: List[ImageLabel]):
        res_img = []
        for image, lab in zip(img_data, labels):
            h, w = image.shape[:2]
            rta_img = rotate_image(image, degree=degree)
            boxes = []
            for obj in lab.boxes:
                minx, miny, maxx, maxy = rotated_calculate_divide_point_bbox((obj.left, obj.top, obj.right, obj.bottom),
                                                                             degree, h, w, number_divide)
                if maxx > minx and maxy > miny:
                    boxes.append(Box(minx, miny, maxx, maxy, obj.name))
            lab.boxes = boxes
            res_img.append(rta_img)
        return res_img, labels

    @staticmethod
    def _save(cvt: Convertor, img_path: str, lbl_path: str, img: np.ndarray, label: ImageLabel):
        label_info = cvt.convert_label(1, label)
        bimg = ImageUtil.encode(img)
        with open(img_path, 'wb') as f:
            f.write(bimg)
        with open(lbl_path, 'wb') as f:
            f.write(label_info.data)

    def save(self, img_data: List[np.ndarray], labels: List[ImageLabel]):
        futures = []
        for img, lbl in zip(img_data, labels):
            if self.contiguous:
                img_name = "{:06}{}".format(self._count, ".jpg")
                lab_name = "{:06}{}".format(self._count, ".txt")
            else:
                img_name = "{:06}{}".format(self._count, ".jpg")
                lab_name = "{:06}{}".format(self._count, ".txt")

            img_path = path_join(self.output, self.convertor.imgDirName, img_name)
            lbl_path = path_join(self.output, self.convertor.lblDirName, lab_name)

            futures.append(self.pro_executor.submit(self._save, self.convertor, img_path, lbl_path, img, lbl))
            self._count += 1


from argparse import ArgumentParser
from datetime import datetime, timedelta, timezone
import logging
import pytz
from tzlocal import get_localzone
from constants import MAX_DAYS, HOURLY_TIME_CHUNK_SIZE

def valid_time(start, end):
    max_historical_time = datetime.now(timezone.utc) - timedelta(days=MAX_DAYS)
    min_time = datetime.now(timezone.utc)
    valid = True
    try:
        datetime(year=start.year, month=start.month, day=start.day, hour=start.hour, minute=start.minute)
    except ValueError:
        logging.error("Invalid start date.")
        valid = False
    except AttributeError:
        logging.error("Invalid attributes for start date.")
        valid = False
    try:
        datetime(year=end.year, month=end.month, day=end.day, hour=end.hour, minute=end.minute)
    except ValueError:
        logging.error("Invalid enddate.")
        valid = False
    except AttributeError:
        logging.error("Invalid attributes for end date.")
        valid = False
    if start.tzinfo is None:
        logging.error("Start time has no timezone. Can't compare")
        valid = False
    elif end.tzinfo is None:
        logging.error("End time has no timezone. Can't compare")
        valid = False
    elif start < max_historical_time:
        logging.error("Start time cannot be more than 30 days in the past.")
        valid = False
    elif start > end + timedelta(days=30):
        logging.error("End time cannot be greater than 30 days after start time.")
        valid = False
    elif end > min_time:
        logging.error("End time cannot be in the future.")
        valid = False
    elif start > end:
        logging.error("Start time cannot be before end time")
        valid = False
    return valid

def split_dates(start=None, end=None):
    # Create a list of tuples (start_time, end_time) that are each chunked CHUNK_SIZE apart + end time
    time_range_list = []
    if valid_time(start, end):
        time_chunk = timedelta(hours=HOURLY_TIME_CHUNK_SIZE)
        time_range_list = []
        while start + time_chunk < end:
            time_range_list.append((start, start + time_chunk))
            start += time_chunk
            logging.debug(f"Append {start} to {start + time_chunk} to time range")
        time_range_list.append((start, end))
        logging.debug(f"Append {start} to {end} to time range")
    logging.debug(f"Split time range into {len(time_range_list)} time chunks.")
    return time_range_list

def add_timezone(time_no_tz, tz=None):
    if time_no_tz.tzinfo in pytz.all_timezones:
        logging.debug("Valid timezone has been provided in ISO string")
        time_tz = time_no_tz
    else:
        logging.debug(f"No timezone in time {time_no_tz} exists.")
        if tz is None:
            local_tz = get_localzone()
            time_tz = time_no_tz.astimezone(local_tz)
            logging.debug(f"Adding system timezone {time_tz.tzinfo} to time {time_no_tz}")
        else:
            try:
                time_tz = time_no_tz.astimezone(pytz.timezone(tz))
                logging.debug(f"Using timezone parameter {time_tz.tzinfo} for time {time_no_tz}")
            except ValueError:
                logging.error(f"Timezone {tz} not a valid time zone.")
                raise ValueError("Not a valid timezone.")

    return time_tz

def check_dates_exist(start, end):
    if end is None:
        end = datetime.now(timezone.utc)
        logging.debug(f"No end date, using today {end}")
    if start is None:
        start = end - timedelta(days=1)
        logging.debug(f"No start date, using today - 1 day {start}")
    return start, end

def convert_timestamp_millisecond(time_convert):
    ms_real = time_convert.timestamp()*1000
    ms_int = round(ms_real)
    return ms_int

def get_date_range(start, end, local_tz=None):
    # Assumes that start_time, end_time are datetime formatted
    (start, end) = check_dates_exist(start, end)
    start_time_tz = add_timezone(start, local_tz)
    end_time_tz = add_timezone(end, local_tz)
    return split_dates(start_time_tz, end_time_tz)

if __name__ == '__main__':
    logging.basicConfig(format='%(levelname)s:%(asctime)s:%(funcName)s():%(message)s',
                        datefmt='%H:%M:%S',
                        level=logging.DEBUG)
    parser = ArgumentParser()
    parser.add_argument("start_time", type=datetime.fromisoformat,
                        help="Start time ISO format [YYY-MM-DDThh:mm:ss.s+TZD]")
    parser.add_argument("end_time", type=datetime.fromisoformat,
                        help="End time ISO format [YYY-MM-DDThh:mm:ss.s+TZD]")
    parser.add_argument("-tz", "--timezone", dest="timezone", type=str,
                        help="Time zone name as per https://en.wikipedia.org/wiki/List_of_tz_database_time_zones")
    args = parser.parse_args()
    time_range = get_date_range(args.start_time, args.end_time, args.timezone)
    for (start_time, end_time) in time_range:
        print("Local",
              start_time.strftime("%Y-%m-%d %H:%M"),
              end_time.strftime("%Y-%m-%d %H:%M"))
        print("UTC",
              start_time.astimezone(timezone.utc).strftime("%Y-%m-%d %H:%M"),
              end_time.astimezone(timezone.utc).strftime("%Y-%m-%d %H:%M"))
        print("Timestamp (msec)",
              convert_timestamp_millisecond(start_time),
              convert_timestamp_millisecond(end_time))


# %%
# load test data
import matplotlib.pyplot as plt
import pickle
import seaborn as sns
import numpy as np
sns.set()
sns.set_context("paper") # notebook, poster,talk, paper
sns.set_style('whitegrid')
sns.palplot(sns.color_palette('hls', 8))
dd = []
with open("eval456.pickle_data", 'rb') as f:
    dd = pickle.load(f)
print(dd)
# %%
# Figure 4
fig, axes = plt.subplots(1, 1, figsize=(8, 4), tight_layout=True)
x = np.arange(1, 29)
x_ticks = np.arange(1, 29)
axes.tick_params(axis='x', labelsize=12,colors='b')
axes.tick_params(axis='y', labelsize=12, colors='b')
axes.bar(x, dd['slew_time'])
axes.set_xlabel('Targets Index',fontsize=14)
axes.set_ylabel('Slew Time (seconds)',fontsize=14)
axes.set_title('')
axes.set_xticks(x_ticks)
axes.grid(True)
plt.savefig("4.pdf", dpi="figure",format="pdf")

# %%
# Figure 5
fig, axes = plt.subplots(1, 1, figsize=(8, 4), tight_layout=True)
x = np.arange(1, 29)
x_ticks = np.arange(1, 29)
axes.tick_params(axis='x', labelsize=12,colors='b')
axes.tick_params(axis='y', labelsize=12, colors='b')
axes.plot(x, dd['start_el'], 's', label="Beginning")
axes.plot(x, dd['end_el'], 'o', label="Ending")
axes.set_xlabel('Targets Index',fontsize=14)
axes.set_ylabel('Elevation Angle(${\degree}$)',fontsize=14)
axes.set_title('')
axes.set_xticks(x_ticks)

axes.legend(loc='upper right')
axes.grid(axis='y')
plt.savefig("5.pdf", dpi="figure",format="pdf")

# %%
# Figure 6
from astropy.time import Time
import astropy.units as u
time_range = Time(['2021-05-21 00:00:00', "2021-05-22 00:00:00"])
time_ut = time_range[0] + np.linspace(0, 24, 24*60)*u.hour

fig, ax = plt.subplots(1, 1, figsize=(12, 4), tight_layout=True)

ax.set_yticks(range(len(dd['target'])))
ax.set_yticklabels(dd['target'])
ax.set_ylim([0, 8])

y=3
text_delt_y = [1, 2, 3, 4, -2 , -3 ,-4]

for i in range(len(dd['target'])-7):
    text_delt_y.append(text_delt_y[len(text_delt_y)-7])

text_delt_x = 0
i = 0
for idx, k in enumerate(dd['target']):
    text_delt_x = 0
    s,e = dd['start_time'][idx],dd['end_time'][idx]
    line1, = ax.plot(range(int(s), int(e)), [y for i in range(int(s), int(e))], '--', linewidth=20,
                    label=k)
    delta_y = text_delt_y[i]
    if k=='0742-2822':
        delta_y = -3
    if delta_y == -4:
        delta_y = 1
        text_delt_x = -100
    ax.annotate(s=k, rotation=18, fontsize=10,xycoords='data',textcoords='data',xy=(int(s)+10, y),
    xytext=(int(s)+text_delt_x, y+delta_y),
    arrowprops=dict(arrowstyle="->",connectionstyle="arc",color="black"),bbox=dict(boxstyle="round,pad=0.1",alpha=0.1,facecolor='green'))
    i = i + 1

ax.set_xlim([0, 1440])
ax.set_yticklabels([])
ax.axes.yaxis.set_visible(False)
xticks = list(range(0, 1400, 100))
#最后一个观测完的时间last_time = 1384
last_time = 1384
xticks.append(last_time)
xticks.append(1440)
ax.set_xticks(xticks)

ax.vlines(last_time, 0, 8, 'gray','--')
plt.setp(ax.get_xticklabels(), rotation=30, ha='right')

d = min(time_ut).datetime
ax.set_xlabel("Minutes offset since {0} {1}(UTC)".format(d.date(), d.time()),fontsize=14)
ax.grid(False)

ax.bar(last_time + (1440-last_time)/2, 8, width=1440-last_time, hatch='//', color='white', edgecolor='gray')

plt.savefig("6.pdf", dpi="figure",format="pdf")
# %%