import mmcv
import os.path as osp
import numpy as np
import sys
from tqdm import tqdm
from detectron2.data import DatasetCatalog, MetadataCatalog
from detectron2.structures import BoxMode
import torch

cur_dir = osp.dirname(osp.abspath(__file__))
sys.path.insert(0, osp.join(cur_dir, "../../../../"))

from lib.vis_utils.colormap import colormap
from lib.utils.mask_utils import cocosegm2mask, get_edge
from core.utils.data_utils import crop_resize_by_warp_affine, read_image_mmcv
from core.gdrn_modeling.datasets.dataset_factory import register_datasets
from lib.pysixd import misc
from transforms3d.quaternions import quat2mat
from lib.egl_renderer.egl_renderer_v3 import EGLRenderer

out_size = 512
score_thr = 0.3
colors = colormap(rgb=False, maximum=255)

id2obj = {
    1: "ape",
    #  2: 'benchvise',
    #  3: 'bowl',
    #  4: 'camera',
    5: "can",
    6: "cat",
    #  7: 'cup',
    8: "driller",
    9: "duck",
    10: "eggbox",
    11: "glue",
    12: "holepuncher",
    #  13: 'iron',
    #  14: 'lamp',
    #  15: 'phone'
}
objects = list(id2obj.values())

tensor_kwargs = {"device": torch.device("cuda"), "dtype": torch.float32}
image_tensor = torch.empty((out_size, out_size, 4), **tensor_kwargs).detach()
seg_tensor = torch.empty((out_size, out_size, 4), **tensor_kwargs).detach()
# image_tensor = torch.empty((480, 640, 4), **tensor_kwargs).detach()

model_dir = "datasets/BOP_DATASETS/lmo/models/"

model_paths = [osp.join(model_dir, f"obj_{obj_id:06d}.ply") for obj_id in id2obj]

ren = EGLRenderer(
    model_paths,
    vertex_scale=0.001,
    use_cache=True,
    width=out_size,
    height=out_size,
)

# NOTE:
pred_path = "output/gdrn/lmo/a6_cPnP_AugAAETrunc_BG0.5_lmo_real_pbr0.1_40e/inference_model_final/lmo_test/a6-cPnP-AugAAETrunc-BG0.5-lmo-real-pbr0.1-40e-test_lmo_test_preds.pkl"

vis_dir = "output/gdrn/lmo/a6_cPnP_AugAAETrunc_BG0.5_lmo_real_pbr0.1_40e/inference_model_final/lmo_test/vis_gt_pred"
mmcv.mkdir_or_exist(vis_dir)

print(pred_path)
preds = mmcv.load(pred_path)

dataset_name = "lmo_test"
print(dataset_name)
register_datasets([dataset_name])

meta = MetadataCatalog.get(dataset_name)
print("MetadataCatalog: ", meta)
objs = meta.objs

dset_dicts = DatasetCatalog.get(dataset_name)
for d in tqdm(dset_dicts):
    K = d["cam"]
    file_name = d["file_name"]
    img = read_image_mmcv(file_name, format="BGR")

    scene_im_id_split = d["scene_im_id"].split("/")
    scene_id = scene_im_id_split[0]
    im_id = int(scene_im_id_split[1])

    imH, imW = img.shape[:2]
    annos = d["annotations"]
    masks = [cocosegm2mask(anno["segmentation"], imH, imW) for anno in annos]
    bboxes = [anno["bbox"] for anno in annos]
    bbox_modes = [anno["bbox_mode"] for anno in annos]
    bboxes_xyxy = np.array(
        [BoxMode.convert(box, box_mode, BoxMode.XYXY_ABS) for box, box_mode in zip(bboxes, bbox_modes)]
    )
    kpts_3d_list = [anno["bbox3d_and_center"] for anno in annos]
    quats = [anno["quat"] for anno in annos]
    transes = [anno["trans"] for anno in annos]
    Rs = [quat2mat(quat) for quat in quats]
    # 0-based label
    cat_ids = [anno["category_id"] for anno in annos]

    kpts_2d = [misc.project_pts(kpt3d, K, R, t) for kpt3d, R, t in zip(kpts_3d_list, Rs, transes)]

    obj_names = [objs[cat_id] for cat_id in cat_ids]

    kpts_2d_est = []
    est_Rs = []
    est_ts = []

    kpts_2d_gt = []
    gt_Rs = []
    gt_ts = []

    kpts_3d_list_sel = []
    labels = []

    maxx, maxy, minx, miny = 0, 0, 1000, 1000
    for anno_i, anno in enumerate(annos):
        kpt_2d_gt = kpts_2d[anno_i]
        obj_name = obj_names[anno_i]

        try:
            R_est = preds[obj_name][file_name]["R"]
            t_est = preds[obj_name][file_name]["t"]
            score = preds[obj_name][file_name]["score"]
        except:
            continue
        if score < score_thr:
            continue

        labels.append(objects.index(obj_name))  # 0-based label

        est_Rs.append(R_est)
        est_ts.append(t_est)

        kpts_3d_list_sel.append(kpts_3d_list[anno_i])
        kpt_2d_est = misc.project_pts(kpts_3d_list[anno_i], K, R_est, t_est)
        kpts_2d_est.append(kpt_2d_est)

        gt_Rs.append(Rs[anno_i])
        gt_ts.append(transes[anno_i])
        kpts_2d_gt.append(kpts_2d[anno_i])

        for i in range(len(kpt_2d_est)):
            maxx, maxy, minx, miny = (
                max(maxx, kpt_2d_est[i][0]),
                max(maxy, kpt_2d_est[i][1]),
                min(minx, kpt_2d_est[i][0]),
                min(miny, kpt_2d_est[i][1]),
            )
            maxx, maxy, minx, miny = (
                max(maxx, kpt_2d_gt[i][0]),
                max(maxy, kpt_2d_gt[i][1]),
                min(minx, kpt_2d_gt[i][0]),
                min(miny, kpt_2d_gt[i][1]),
            )
    center = np.array([(minx + maxx) / 2, (miny + maxy) / 2])
    scale = max(maxx - minx, maxy - miny) * 1.5  # + 10
    crop_minx = max(0, center[0] - scale / 2)
    crop_miny = max(0, center[1] - scale / 2)
    crop_maxx = min(imW - 1, center[0] + scale / 2)
    crop_maxy = min(imH - 1, center[1] + scale / 2)
    scale = min(scale, min(crop_maxx - crop_minx, crop_maxy - crop_miny))

    zoomed_im = crop_resize_by_warp_affine(img, center, scale, out_size)
    im_zoom_gray = mmcv.bgr2gray(zoomed_im, keepdim=True)
    im_zoom_gray_3 = np.concatenate([im_zoom_gray, im_zoom_gray, im_zoom_gray], axis=2)
    # print(im_zoom_gray.shape)
    K_zoom = K.copy()
    K_zoom[0, 2] -= center[0] - scale / 2
    K_zoom[1, 2] -= center[1] - scale / 2
    K_zoom[0, :] *= out_size / scale
    K_zoom[1, :] *= out_size / scale

    gt_poses = [np.hstack([_R, _t.reshape(3, 1)]) for _R, _t in zip(gt_Rs, gt_ts)]
    poses = [np.hstack([_R, _t.reshape(3, 1)]) for _R, _t in zip(est_Rs, est_ts)]

    ren.render(
        labels,
        poses,
        K=K_zoom,
        image_tensor=image_tensor,
        background=im_zoom_gray_3,
    )
    ren_bgr = (image_tensor[:, :, :3].detach().cpu().numpy() + 0.5).astype("uint8")

    # gt_masks = []
    # est_masks = []
    for label, gt_pose, est_pose in zip(labels, gt_poses, poses):
        ren.render([label], [gt_pose], K=K_zoom, seg_tensor=seg_tensor)
        gt_mask = (seg_tensor[:, :, 0].detach().cpu().numpy() > 0).astype("uint8")

        ren.render([label], [est_pose], K=K_zoom, seg_tensor=seg_tensor)
        est_mask = (seg_tensor[:, :, 0].detach().cpu().numpy() > 0).astype("uint8")

        gt_edge = get_edge(gt_mask, bw=3, out_channel=1)
        est_edge = get_edge(est_mask, bw=3, out_channel=1)

        # zoomed_im[gt_edge != 0] = np.array(mmcv.color_val("blue"))
        # zoomed_im[est_edge != 0] = np.array(mmcv.color_val("green"))

        ren_bgr[gt_edge != 0] = np.array(mmcv.color_val("blue"))
        ren_bgr[est_edge != 0] = np.array(mmcv.color_val("green"))

    vis_im = ren_bgr

    # vis_im_add = (im_zoom_gray_3 * 0.3 + ren_bgr * 0.7).astype("uint8")

    # kpts_2d_gt_zoom = [misc.project_pts(kpt3d, K_zoom, R, t) for kpt3d, R, t in zip(kpts_3d_list_sel, gt_Rs, gt_ts)]
    # kpts_2d_est_zoom = [misc.project_pts(kpt3d, K_zoom, R, t) for kpt3d, R, t in zip(kpts_3d_list_sel, est_Rs, est_ts)]
    # linewidth = 3
    # visualizer = MyVisualizer(zoomed_im[:, :, ::-1], meta)
    # for kpt_2d_gt_zoom, kpt_2d_est_zoom in zip(kpts_2d_gt_zoom, kpts_2d_est_zoom):
    #     visualizer.draw_bbox3d_and_center(
    #         kpt_2d_gt_zoom, top_color=_BLUE, bottom_color=_GREY, linewidth=linewidth, draw_center=True
    #     )
    #     visualizer.draw_bbox3d_and_center(
    #         kpt_2d_est_zoom, top_color=_GREEN, bottom_color=_GREY, linewidth=linewidth, draw_center=True
    #     )
    # vis_im = visualizer.get_output()
    # save_path = osp.join(vis_dir, "{}_{:06d}_gt_est.png".format(scene_id, im_id))
    # vis_im.save(save_path)

    save_path_0 = osp.join(vis_dir, "{}_{:06d}_im.png".format(scene_id, im_id))
    mmcv.imwrite(zoomed_im, save_path_0)

    save_path = osp.join(vis_dir, "{}_{:06d}_gt_est.png".format(scene_id, im_id))
    mmcv.imwrite(vis_im, save_path)

    # if True:
    #     # grid_show([zoomed_im[:, :, ::-1], vis_im[:, :, ::-1]], ["im", "est"], row=1, col=2)
    #     # im_show = cv2.hconcat([zoomed_im, vis_im, vis_im_add])
    #     im_show = cv2.hconcat([zoomed_im, vis_im])
    #     cv2.imshow("im_est", im_show)
    #     if cv2.waitKey(0) == 27:
    #         break  # esc to quit


import pygame
import sys
class Main:
    def __init__(self, scr_size):
        pygame.init()
        pygame.freetype.init()
        self.size = scr_size
        self.screen = pygame.display.set_mode(self.size)
        self.clock = pygame.time.Clock()
        self.process_stopped = False
        self.score = 0

    def quit(self, e):
        if e.type == pygame.QUIT:
            self.process_stopped = True
            sys.exit(0)
        if e.type == pygame.KEYDOWN:
            if e.key == pygame.K_ESCAPE:
                self.process_stopped = True
                sys.exit(0)


from django.conf import settings

DEFAULT_SETTINGS = {
    'ALLOW_VERSION_OVERWRITE': False,
    'RELEASE_UPLOAD_TO': 'dists',
    'OS_NAMES': (
        ("aix", "AIX"),
        ("beos", "BeOS"),
        ("debian", "Debian Linux"),
        ("dos", "DOS"),
        ("freebsd", "FreeBSD"),
        ("hpux", "HP/UX"),
        ("mac", "Mac System x."),
        ("macos", "MacOS X"),
        ("mandrake", "Mandrake Linux"),
        ("netbsd", "NetBSD"),
        ("openbsd", "OpenBSD"),
        ("qnx", "QNX"),
        ("redhat", "RedHat Linux"),
        ("solaris", "SUN Solaris"),
        ("suse", "SuSE Linux"),
        ("yellowdog", "Yellow Dog Linux"),
    ),
    'ARCHITECTURES': (
        ("alpha", "Alpha"),
        ("hppa", "HPPA"),
        ("ix86", "Intel"),
        ("powerpc", "PowerPC"),
        ("sparc", "Sparc"),
        ("ultrasparc", "UltraSparc"),
    ),
    'DIST_FILE_TYPES': (
        ('sdist', 'Source'),
        ('bdist_dumb', '"dumb" binary'),
        ('bdist_rpm', 'RPM'),
        ('bdist_wininst', 'MS Windows installer'),
        ('bdist_egg', 'Python Egg'),
        ('bdist_dmg', 'OS X Disk Image'),
    ),
    'PYTHON_VERSIONS': (
        ('any', 'Any i.e. pure python'),
        ('2.1', '2.1'),
        ('2.2', '2.2'),
        ('2.3', '2.3'),
        ('2.4', '2.4'),
        ('2.5', '2.5'),
        ('2.6', '2.6'),
        ('2.7', '2.7'),
        ('3.0', '3.0'),
        ('3.1', '3.1'),
        ('3.2', '3.2'),
    ),
    'METADATA_FIELDS': {
        '1.0': ('platform', 'summary', 'description', 'keywords', 'home_page',
                'author', 'author_email', 'license', ),
        '1.1': ('platform', 'supported_platform', 'summary', 'description',
                'keywords', 'home_page', 'download_url', 'author', 'author_email',
                'license', 'classifier', 'requires', 'provides', 'obsoletes',),
        '1.2': ('platform', 'supported_platform', 'summary', 'description',
                'keywords', 'home_page', 'download_url', 'author', 'author_email',
                'maintainer', 'maintainer_email', 'license', 'classifier',
                'requires_dist', 'provides_dist', 'obsoletes_dist',
                'requires_python', 'requires_external', 'project_url')
    },
    'METADATA_FORMS': {
        '1.0': 'djangopypi.forms.Metadata10Form',
        '1.1': 'djangopypi.forms.Metadata11Form',
        '1.2': 'djangopypi.forms.Metadata12Form',
    },
    'FALLBACK_VIEW': 'djangopypi.views.releases.ReleaseListView',
    'ACTION_VIEWS': {
        "file_upload": 'djangopypi.views.distutils.register_or_upload',  # ``sdist`` command
        "submit": 'djangopypi.views.distutils.register_or_upload',  # ``register`` command
        "list_classifiers": 'djangopypi.views.distutils.list_classifiers',  # ``list_classifiers`` command
    },
    'XMLRPC_COMMANDS': {
        'list_packages': 'djangopypi.views.xmlrpc.list_packages',
        'package_releases': 'djangopypi.views.xmlrpc.package_releases',
        'release_urls': 'djangopypi.views.xmlrpc.release_urls',
        'release_data': 'djangopypi.views.xmlrpc.release_data',
        # 'search': xmlrpc.search, Not done yet
        # 'changelog': xmlrpc.changelog, Not done yet
        # 'ratings': xmlrpc.ratings, Not done yet
    },
    'MIRRORING': False,
    'PROXY_BASE_URL': 'http://pypi.python.org',
    'PROXY_MISSING': False,
}

USER_SETTINGS = DEFAULT_SETTINGS.copy()
USER_SETTINGS.update(getattr(settings, 'DJANGOPYPI_SETTINGS', {}))

# This is disabled on pypi.python.org, can be useful if you make mistakes
if hasattr(settings, 'DJANGOPYPI_ALLOW_VERSION_OVERWRITE'):
    USER_SETTINGS['ALLOW_VERSION_OVERWRITE'] = settings.DJANGOPYPI_ALLOW_VERSION_OVERWRITE

""" The upload_to argument for the file field in releases. This can either be
a string for a path relative to your media folder or a callable. For more
information, see http://docs.djangoproject.com/ """
if hasattr(settings, 'DJANGOPYPI_RELEASE_UPLOAD_TO'):
    USER_SETTINGS['RELEASE_UPLOAD_TO'] = settings.DJANGOPYPI_RELEASE_UPLOAD_TO

if hasattr(settings, 'DJANGOPYPI_OS_NAMES'):
    USER_SETTINGS['OS_NAMES'] = settings.DJANGOPYPI_OS_NAMES

if hasattr(settings, 'DJANGOPYPI_ARCHITECTURES'):
    USER_SETTINGS['ARCHITECTURES'] = settings.DJANGOPYPI_ARCHITECTURES

if hasattr(settings, 'DJANGOPYPI_DIST_FILE_TYPES'):
    USER_SETTINGS['DIST_FILE_TYPES'] = settings.DJANGOPYPI_DIST_FILE_TYPES

if hasattr(settings, 'DJANGOPYPI_PYTHON_VERSIONS'):
    USER_SETTINGS['PYTHON_VERSIONS'] = settings.DJANGOPYPI_PYTHON_VERSIONS

if hasattr(settings, 'DJANGOPYPI_METADATA_FIELDS'):
    USER_SETTINGS['METADATA_FIELDS'] = settings.DJANGOPYPI_METADATA_FIELDS

if hasattr(settings, 'DJANGOPYPI_METADATA_FORMS'):
    USER_SETTINGS['METADATA_FORMS'] = settings.DJANGOPYPI_METADATA_FORMS

if hasattr(settings, 'DJANGOPYPI_FALLBACK_VIEW'):
    USER_SETTINGS['FALLBACK_VIEW'] = settings.DJANGOPYPI_FALLBACK_VIEW

if hasattr(settings, 'DJANGOPYPI_ACTION_VIEWS'):
    USER_SETTINGS['ACTION_VIEWS'] = settings.DJANGOPYPI_ACTION_VIEWS

if hasattr(settings, 'DJANGOPYPI_XMLRPC_COMMANDS'):
    USER_SETTINGS['XMLRPC_COMMANDS'] = settings.DJANGOPYPI_XMLRPC_COMMANDS

""" These settings enable proxying of packages that are not in the local index
to another index, http://pypi.python.org/ by default. This feature is disabled
by default and can be enabled by setting DJANGOPYPI_PROXY_MISSING to True in
your settings file. """
if hasattr(settings, 'DJANGOPYPI_PROXY_BASE_URL'):
    USER_SETTINGS['PROXY_BASE_URL'] = settings.DJANGOPYPI_PROXY_BASE_URL

if not hasattr(settings, 'DJANGOPYPI_PROXY_MISSING'):
    USER_SETTINGS['PROXY_MISSING'] = settings.DJANGOPYPI_PROXY_MISSING = False

# USER_SETTINGS['FALLBACK_VIEW'] = import_item(USER_SETTINGS['FALLBACK_VIEW'])
# USER_SETTINGS['METADATA_FORMS'] = import_item(USER_SETTINGS['METADATA_FORMS'])
# action_views_dict = {}
# for key, val in USER_SETTINGS['ACTION_VIEWS'].items():
#     action_views_dict[key] = import_item(USER_SETTINGS['ACTION_VIEWS'][key])
# USER_SETTINGS['ACTION_VIEWS'] = action_views_dict
# xmlrpc_commands = {}
# for key, val in USER_SETTINGS['XMLRPC_COMMANDS'].items():
#     xmlrpc_commands[key] = import_item(USER_SETTINGS['XMLRPC_COMMANDS'][key])
# USER_SETTINGS['XMLRPC_COMMANDS'] = xmlrpc_commands

globals().update(USER_SETTINGS)


import paddle
from efficientnet import EfficientNet

def EfficientNetmy(padding_type='SAME',
                   override_params=None,
                   use_se=True,
                   pretrained=False,
                   use_ssld=False,
                   **kwargs):
    model = EfficientNet(
        name='my',
        padding_type=padding_type,
        override_params=override_params,
        use_se=use_se,
        **kwargs)
    # _load_pretrained(pretrained, model, MODEL_URLS["EfficientNetB0"])
    return model

model = EfficientNetmy()
layer_state_dict = model.state_dict()
for k in layer_state_dict:
    print(k)


from django.db.models import Q
from django.db.models.functions import Lower
from django.urls import reverse_lazy
from poll.models import Poll
from accomplishment.serializers import AccomplishmentSerializer
from uniklinik.mixins import DatatablesMixin

class PollDatatables(DatatablesMixin):
    queryset = Poll.objects.all()
    serializer_class = AccomplishmentSerializer

    def __init__(self):
        super().__init__()
        self.records_total = None
        self.page_number = None
        self.page_size = None

    def get_filtered_queryset(self):
        self.filter_by_search_value()
        return self.queryset

    def filter_by_search_value(self):
        whole_search_value = self.request.GET.get("search[value]")
        search_values = whole_search_value.split(" ")
        for search_value in search_values:
            if search_value != "" and search_value is not None:
                self.queryset = self.queryset.filter(
                    Q(Q(title__icontains=search_value))).distinct()

    def get_ordered_queryset(self):
        order_column_index = self.request.GET.get("order[0][column]")
        asc_or_desc = self.request.GET.get("order[0][dir]")
        if order_column_index == "1":
            if asc_or_desc == "asc":
                self.queryset = self.queryset.order_by(Lower("title"))
            else:
                self.queryset = self.queryset.annotate(lower_title=Lower('title')).order_by("-lower_title")
        return self.queryset

    def get_data(self, page):
        data = {"results": [], "records_total": self.queryset.count()}
        for query in page:
            if query.open is True:
                poll_is_open_or_not = f''
            else:
                poll_is_open_or_not = f''

            data["results"].append([
                f'Bearbeiten',
                # f'Bearbeiten',
                query.title, poll_is_open_or_not])
        return data


import numpy as np

from garage.np.baselines.base import Baseline

class LinearFeatureBaseline(Baseline):

    def __init__(self, env_spec, reg_coeff=1e-5, name='LinearFeatureBaseline'):
        self._coeffs = None
        self._reg_coeff = reg_coeff
        self.name = name

    def get_param_values(self, **tags):
        return self._coeffs

    def set_param_values(self, val, **tags):
        self._coeffs = val

    def _features(self, path):
        obs = np.clip(path['observations'], -10, 10)
        length = len(path['rewards'])
        al = np.arange(length).reshape(-1, 1) / 100.0
        return np.concatenate(
            [obs, obs**2, al, al**2, al**3,
             np.ones((length, 1))], axis=1)

    def fit(self, paths):
        featmat = np.concatenate([self._features(path) for path in paths])
        returns = np.concatenate([path['returns'] for path in paths])
        reg_coeff = self._reg_coeff
        for _ in range(5):
            self._coeffs = np.linalg.lstsq(
                featmat.T.dot(featmat) +
                reg_coeff * np.identity(featmat.shape[1]),
                featmat.T.dot(returns),
                rcond=-1)[0]
            if not np.any(np.isnan(self._coeffs)):
                break
            reg_coeff *= 10

    def predict(self, path):
        if self._coeffs is None:
            return np.zeros(len(path['rewards']))
        return self._features(path).dot(self._coeffs)


from sqlalchemy import orm

from ..core.qt import QtGui
from ..admin.action.list_action import ListActionGuiContext

class MockModelContext( object ):
    """Model Context to be used in unit tests

    :param session: the session attributed to this model context, if `None` is
        given, the session of the object is used.
    """

    def __init__( self, session=None ):
        self._model = []
        self.obj = None
        self.selection = []
        self.admin = None
        self.mode_name = None
        self.collection_count = 1
        self.selection_count = 1
        self.current_row = 0
        self.current_column = None
        self.current_field_name = None
        self.field_attributes = {}
        self._session = session

    def get_object( self ):
        return self.obj

    def get_selection( self, yield_per = None ):
        if self.obj is not None:
            return [self.obj]
        return self.selection

    def get_collection( self, yield_per = None ):
        return self.get_selection(yield_per=yield_per)

    @property
    def session( self ):
        return self._session or orm.object_session( self.obj )

class MockListActionGuiContext( ListActionGuiContext ):

    def __init__( self ):
        super(MockListActionGuiContext, self).__init__()
        self.item_view = QtGui.QTableWidget( 4, 4 )


import sys, os, argparse
import scad
from typing import List

def do_inline(src, dest, scad_paths):
    try:
        scad.inliner(src, dest, scad_paths)
    except FileNotFoundError as e:
        print(e)
        print('Try to use the -op argument or set up the OPENSCADPATH environment variable')

def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("src", help="the source")
    parser.add_argument("dest", help="the destination")
    parser.add_argument("-op", "--openscadpath", help="specify where to find user scad files")

    if(len(sys.argv) == 1):
        parser.print_help()
    else:
        args = parser.parse_args()

        scad_paths: List[str] = (
            args.openscadpath
                if args.openscadpath else os.environ['OPENSCADPATH']
        ).split(';')

        do_inline(args.src, args.dest, scad_paths)

main()


import argparse
import os

import torch
from torch.utils.data import DataLoader
from egg import core
import neptune
from neptunecontrib.api.utils import get_filepaths

from compositionality.games import PretrainingmGameReinforce, PretrainingmGameGS, CompositionalGameGS, CompositionalGameReinforce
from compositionality.wrappers import RnnReceiverDeterministic, RnnReceiverGS
from compositionality.callbacks import NeptuneMonitor, EarlyStopperAccuracy, CompositionalityMetricReinforce, CompositionalityMetricGS
from compositionality.agents import Sender, Receiver, VisualSender
from compositionality.data import prepare_datasets

from visual_compositionality.visual_data import prepare_datasets

def get_params():
    parser = argparse.ArgumentParser()
    parser.add_argument('--n_features', type=int, default=5,
                        help='Dimensionality of the "concept" space (default: 10)')
    parser.add_argument('--n_attributes', type=int, default=2,
                        help='Number of attributes (default: 2')
    parser.add_argument('--sender_hidden', type=int, default=200,
                        help='Size of the hidden layer of Sender (default: 200)')
    parser.add_argument('--receiver_hidden', type=int, default=200,
                        help='Size of the hidden layer of Receiver (default: 200)')
    parser.add_argument('--sender_embedding', type=int, default=50,
                        help='Dimensionality of the embedding hidden layer for Sender (default: 10)')
    parser.add_argument('--receiver_embedding', type=int, default=50,
                        help='Dimensionality of the embedding hidden layer for Receiver (default: 10)')
    parser.add_argument('--rnn_cell', type=str, default='rnn')
    parser.add_argument('--sender_lr', type=float, default=0.001,
                        help="Learning rate for Sender's parameters (default: 1e-3)")
    parser.add_argument('--receiver_lr', type=float, default=0.001,
                        help="Learning rate for Receiver's parameters (default: 1e-3)")
    parser.add_argument('--seed', type=int, default=171,
                        help="Random seed")
    parser.add_argument('--pretrain', action='store_true', default=False, help="")
    parser.add_argument('--reinforce', action='store_true', default=False, help="")
    parser.add_argument('--config', type=str, default=None)

    args = core.init(parser)
    print(args)
    return args

if __name__ == "__main__":
    opts = get_params()
    opts.on_slurm = os.environ.get('SLURM_JOB_NAME', False)
    core.util._set_seed(opts.seed)
    full_dataset, train, test = prepare_datasets(opts.n_features, opts.n_attributes)
    train_loader = DataLoader(train, batch_size=opts.batch_size, drop_last=False, shuffle=True)
    test_loader = DataLoader(test, batch_size=opts.batch_size, drop_last=False, shuffle=False)

    if opts.reinforce:
        pretrained_senders = [
            core.RnnSenderReinforce(
                agent=Sender(opts.sender_hidden, opts.n_features, opts.n_attributes),
                vocab_size=opts.vocab_size,
                embed_dim=opts.sender_embedding,
                hidden_size=opts.sender_hidden,
                max_len=1,
                cell=opts.rnn_cell,
                force_eos=False
            )
            for i in range(2)]
        sender_3 = core.RnnSenderReinforce(
            agent=Sender(opts.sender_hidden, opts.n_features, opts.n_attributes),
            vocab_size=opts.vocab_size,
            embed_dim=opts.sender_embedding,
            hidden_size=opts.sender_hidden,
            max_len=opts.max_len,
            force_eos=False,
            cell=opts.rnn_cell)
        receiver = RnnReceiverDeterministic(
            agent=Receiver(opts.receiver_hidden, opts.n_features, opts.n_attributes),
            vocab_size=opts.vocab_size,
            embed_dim=opts.receiver_embedding,
            hidden_size=opts.receiver_hidden,
            cell=opts.rnn_cell)
    else:
        pretrained_senders = [
            core.RnnSenderGS(
                agent=VisualSender(opts.sender_hidden, opts.n_features, opts.n_attributes),
                vocab_size=opts.vocab_size,
                embed_dim=opts.sender_embedding,
                hidden_size=opts.sender_hidden,
                max_len=1,
                temperature=3.,
                trainable_temperature=True,
                cell=opts.rnn_cell,
                force_eos=False
            )
            for i in range(2)]
        sender_3 = core.RnnSenderGS(
                agent=VisualSender(opts.sender_hidden, opts.n_features, opts.n_attributes),
                vocab_size=opts.vocab_size,
                embed_dim=opts.sender_embedding,
                hidden_size=opts.sender_hidden,
                max_len=opts.max_len,
                temperature=3.,
                trainable_temperature=True,
                force_eos=False,
                cell=opts.rnn_cell)
        receiver = RnnReceiverGS(
                agent=Receiver(opts.receiver_hidden, opts.n_features, opts.n_attributes),
                vocab_size=opts.vocab_size,
                embed_dim=opts.receiver_embedding,
                hidden_size=opts.receiver_hidden,
                cell=opts.rnn_cell)

    neptune.init('tomekkorbak/template-transfer')
    with neptune.create_experiment(params=vars(opts), upload_source_files=get_filepaths(), tags=['rrr']) as experiment:
        CompositionalityMetric = CompositionalityMetricReinforce if opts.reinforce else CompositionalityMetricGS

        # Pretraining game
        if opts.pretrain:
            pretraining_game = PretrainingmGameReinforce(pretrained_senders, receiver) if opts.reinforce else PretrainingmGameGS(pretrained_senders, receiver)
            sender_params = [{'params': sender.parameters(), 'lr': opts.sender_lr} for sender in pretrained_senders]
            receiver_params = [{'params': receiver.parameters(), 'lr': opts.receiver_lr}]
            optimizer = torch.optim.Adam(sender_params + receiver_params)
            trainer = core.Trainer(
                game=pretraining_game, optimizer=optimizer, train_data=train_loader,
                validation_data=test_loader,
                callbacks=[
                    CompositionalityMetric(full_dataset, pretrained_senders[0], opts, opts.vocab_size, prefix='1_'),
                    CompositionalityMetric(full_dataset, pretrained_senders[1], opts, opts.vocab_size, prefix='2_'),
                    NeptuneMonitor(prefix='pretrain'),
                    core.ConsoleLogger(print_train_loss=not opts.on_slurm),
                    EarlyStopperAccuracy(threshold=0.85, field_name='accuracy', delay=2),
                ])
            trainer.train(n_epochs=500_000)
            pretraining_game.train(False)

        # Compositional game
        compositional_game = CompositionalGameReinforce(sender_3, receiver) if opts.reinforce else CompositionalGameGS(sender_3, receiver)
        sender_params = [{'params': sender_3.parameters(), 'lr': opts.sender_lr}]
        receiver_params = [{'params': receiver.parameters(), 'lr': opts.receiver_lr * 0.5}]
        optimizer = torch.optim.Adam(sender_params + receiver_params)
        trainer = core.Trainer(game=compositional_game, optimizer=optimizer, train_data=train_loader,
                               validation_data=test_loader,
                               callbacks=[
                                   CompositionalityMetric(full_dataset, sender_3, opts, opts.vocab_size, prefix='comp'),
                                   NeptuneMonitor(prefix='comp'),
                                   core.ConsoleLogger(print_train_loss=not opts.on_slurm),
                               ])
        trainer.train(n_epochs=opts.n_epochs)


import setuptools

with open("README.md", "r") as fh:
    long_description = fh.read()

print(setuptools.find_packages())

setuptools.setup(
    name="TerraVide",                     # This is the name of the package
    version="0.1.7",                        #release version
    author="Sarang Pramode",                     # Full name of the author
    description="An open source python package to process and simulate large urban environments mapped with LiDAR data",
    long_description = long_description,
    long_description_content_type = "text/markdown",
    license='MIT',
    packages=setuptools.find_packages(),    # List of all python modules to be installed
    classifiers=[
        "Programming Language :: Python :: 3",
        "License :: OSI Approved :: MIT License",
        "Operating System :: OS Independent",
    ],                                      # Information to filter the project on PyPi website
    python_requires='>=3.6',                # Minimum version requirement of the package
    #py_modules=["info","dataset"],             # Name of the python package
    install_requires=[              # Install other dependencies if any
        "numpy",
        "pandas",
        "laspy",
    ]
)


import os

import affine
from tile2net.logger import logger
import shapely
from geopandas import GeoDataFrame, read_file, sjoin
import pandas as pd
import numpy as np
import rasterio
from shapely.geometry import mapping, shape
import skimage
from affine import Affine

import warnings

warnings.filterwarnings('ignore')

def read_gdf(path):
    """
    Read a GeoDataFrame from a file
    Parameters
    ----------
    path: str
        path to the file

    Returns
    -------
    gdf: GeoDataFrame
    """
    gdf = read_file(path)
    return gdf

def set_gdf_crs(gdf, crs):
    """
    Set the CRS of a GeoDataFrame
    Parameters
    ----------
    gdf: GeoDataFrame
    crs: int
        coordinate reference system

    Returns
    -------
    gdf: GeoDataFrame
    """
    gdf.geometry = gdf.geometry.set_crs(crs)
    return gdf

def change_crs(gdf, crs):
    """
    Change the CRS of a GeoDataFrame
    Parameters
    ----------
    gdf: GeoDataFrame
    crs: int

    Returns
    -------
    gdf: GeoDataFrame
    """
    gdf.geometry = gdf.geometry.to_crs(crs)
    return gdf

def prepare_spindex(gdf: GeoDataFrame) -> GeoDataFrame.sindex:
    """
    Prepare a GeoDataFrame for spatial indexing
    Parameters
    ----------
    gdf: GeoDataFrame

    Returns
    -------
    spatial index of a GeoDataFrame
    """
    return gdf.sindex

def _reduce_geom_precision(geom, precision=2):
    """
    Reduce the precision of a geometry to a given number of decimal places.
    Parameters
    ----------
    geom: shapely.geometry
    precision: int
        number of decimal places to round to

    Returns
    -------
    geom: shapely.geometry
    """
    geojson = mapping(geom)
    geojson['coordinates'] = np.round(
        np.array(geojson['coordinates']),
        precision
    )
    return shape(geojson)

def affine_to_list(affine_obj: affine.Affine):
    """Convert a :class:`affine.Affine` instance to a list for Shapely."""
    return [affine_obj.a, affine_obj.b,
            affine_obj.d, affine_obj.e,
            affine_obj.xoff, affine_obj.yoff]

def list_to_affine(xform_mat: list):
    """Create an Affine from a list or array-formatted [a, b, d, e, xoff, yoff]

    Arguments
    ---------
    xform_mat : `list` or :class:`numpy.array`
        A `list` of values to convert to an affine object.

    Returns
    -------
    aff : :class:`affine.Affine`
        An affine transformation object.
    """
    # first make sure it's not in gdal order
    if len(xform_mat) > 6:
        xform_mat = xform_mat[0:6]
    if rasterio.transform.tastes_like_gdal(xform_mat):
        return Affine.from_gdal(*xform_mat)
    else:
        return Affine(*xform_mat)

def _check_rasterio_im_load(im):
    """Check if `im` is already loaded in; if not, load it in."""
    if isinstance(im, str):
        return rasterio.open(im)
    elif isinstance(im, rasterio.DatasetReader):
        return im
    else:
        raise ValueError(
            "{} is not an accepted image format for rasterio.".format(im)
        )

def _check_skimage_im_load(im):
    """Check if `im` is already loaded in; if not, load it in."""
    if isinstance(im, str):
        return skimage.io.imread(im)
    elif isinstance(im, np.ndarray):
        return im
    else:
        raise ValueError(
            "{} is not an accepted image format for scikit-image.".format(im)
        )

def prepare_class_gdf(polys, class_name) -> GeoDataFrame:
    """
    separates the polygons of each class, given the keyboard (sidewalk, crosswalk, road)
    Args:
        polys (geodataframe): the dataframe containing the polygons of all classes
        class_name(str): the class label, sidewalk, crosswalk, road

    Returns:
        class specific GeoDataFrame in metric projection
    """

    nt = polys[polys.f_type == f'{class_name}'].copy()
    nt.geometry = nt.geometry.to_crs(3857)
    return nt

def prepare_gdf(gdf, **cols):
    """
    Filter a GeoDataFrame based on a set of columns and values
    Parameters
    ----------
    gdf: GeoDataFrame
    cols: dict
        {column_name: value}

    Returns
    -------
    f_gdf: GeoDataFrame
    """
    # TODO: Add other operations like !
    k = list(cols.keys())[0]
    print(f'k, {k}', f'cols {len(cols[k])}')
    if isinstance(cols[k], list):
        f_gdf = gdf[gdf[k].isin(cols[k])]
    else:
        f_gdf = gdf[gdf[k] == cols[k]]
    return f_gdf

def read_dataframe(src_path, geo=True, cols=None):
    """
    Args:
        src_path:
        geo: if True, will create GeoDataFrame
        cols: optional. Name of specific columns to be read
    Returns:
    """
    if geo:
        if cols:
            df = read_file(src_path, usecols=cols)
        else:
            df = read_file(src_path)
    else:
        if cols:
            df = pd.read_csv(src_path, usecols=cols)
        else:
            df = pd.DataFrame(src_path)
    return df

def unary_multi(gdf: GeoDataFrame) -> GeoDataFrame:
    # handles the errors with multipolygon
    loc = ~gdf.is_valid.values
    # logger.warning(f'Number of invalid geometries: {loc.sum()} out of {len(gdf)}')
    count = loc.sum()
    if count:
        logger.warning(f'Number of invalid geometries: {count} out of {len(gdf)}')
    gdf.geometry.loc[loc] = shapely.make_valid(gdf.geometry.loc[loc])
    result = (
        gdf
        # dissolve overlapping geometries
        .dissolve()
        # explode multipart geometries
        .explode()
    )
    return result

def buffer_union(gdf, buff, simp1, simp2):
    """
    buffer and union the polygons in a GeoDataFrame
    Parameters
    ----------
    gdf: GeoDataFrame
    buff: float
        buffer distance
    simp1: float
        simplification tolerance for the buffer
    simp2: float
        simplification tolerance for the union

    Returns
    -------
    gdf_uni: GeoDataFrame
    """
    gdf.geometry = gdf.geometry.buffer(buff, join_style=2, cap_style=3)
    gdf.geometry = gdf.simplify(simp1)
    gdf_uni = unary_multi(gdf)
    gdf_uni.geometry = gdf_uni.geometry.set_crs(3857)
    gdf_uni.geometry = gdf_uni.geometry.simplify(simp2)
    return gdf_uni

def buffer_union_erode(gdf: GeoDataFrame, buff: float, erode: float, simp1: float, simp2: float, simp3: float):
    """ Buffer, union, erode, simplify the polygons in a GeoDataFrame to create elongated polygons
    Parameters
    ----------
    gdf: GeoDataFrame
        Polygon GeoDataFrame
    buff: float
        buffer distance
    erode: float
        erode distance
    simp1: float
        simplification tolerance for the buffer
    simp2: float
        simplification tolerance for the union
    simp3: float
        simplification tolerance for the final simplification
    Returns
    -------
    gdf_uni: GeoDataFrame
        Transformed GeoDataFrame
    """
    gdf_buff = buffer_union(gdf, buff, simp1, simp2)
    gdf_erode = gdf_buff.copy()
    gdf_erode.geometry = gdf_buff.geometry.buffer(erode, join_style=2, cap_style=3)
    gdf_uni = unary_multi(gdf_erode)
    gdf_uni.geometry = gdf_uni.geometry.set_crs(3857)
    gdf_uni.geometry = gdf_uni.geometry.simplify(simp3)
    gdf_uni = gdf_uni[(gdf_uni.geometry.geom_type == 'MultiPolygon') | (gdf_uni.geometry.geom_type == 'Polygon')]

    return gdf_uni

def to_metric(gdf, crs=3857):
    """Converts a GeoDataFrame to metric (3857) coordinate
    Parameters
    ----------
    gdf : GeoDataFrame
        GeoDataFrame of polygons
    crs : int, optional
        the coordinate system to convert to, by default 3857
    Returns
    -------
    GeoDataFrame
        GeoDataFrame of polygons in metric coordinate system
    """
    gdf.to_crs(crs, inplace=True)
    return gdf

def geo2geodf(geo_lst: list) -> GeoDataFrame:
    """
    Converts a list of shapely geometries to a GeoDataFrame
    Parameters
    ----------
    geo_lst: list

    Returns
    -------
    gdf: GeoDataFrame

    """
    gdf = GeoDataFrame(geometry=geo_lst)
    return gdf

def merge_dfs(gdf1: GeoDataFrame, gdf2: GeoDataFrame, crs: int = 4326):
    """
    merges two geodataframes of segmentation results (three classes)
    Parameters
    ----------
    gdf1: GeoDataFrame
    gdf2: GeoDataFrame
    crs: int
        Coordinate system to convert to
    Returns
    -------
    concsw: GeoDataFrame
        Concatenated GeoDataFrames
    """
    if gdf1.crs != gdf2.crs:
        gdf1.to_crs(crs, inplace=True)
        gdf2.to_crs(crs, inplace=True)

    df1sw = prepare_class_gdf(gdf1, 'sidewalk')
    df1cw = prepare_class_gdf(gdf1, 'crosswalk')
    df1rd = prepare_class_gdf(gdf1, 'road')

    df2sw = prepare_class_gdf(gdf2, 'sidewalk')
    df2cw = prepare_class_gdf(gdf2, 'crosswalk')
    df2rd = prepare_class_gdf(gdf2, 'road')

    concsw = pd.concat([df1sw, df2sw])
    conccw = pd.concat([df1cw, df2cw])
    concrd = pd.concat([df1rd, df2rd])

    unionsw = unary_multi(concsw)
    unionsw = unionsw.explode().reset_index(drop=True)
    unionsw.geometry = unionsw.geometry.set_crs(crs)
    unionsw['f_type'] = 'sidewalk'

    unioncw = unary_multi(conccw)

    unioncw.geometry = unioncw.geometry.set_crs(crs)
    unioncw['f_type'] = 'crosswalk'

    unionrd = unary_multi(concrd)

    unionrd.geometry = unionrd.geometry.set_crs(crs)
    unionrd['f_type'] = 'road'

    merged = pd.concat([unionrd, unionsw, unioncw])
    merged.geometry = merged.geometry.set_crs(crs)

    return merged

def create_stats(gdf: GeoDataFrame):
    """ Creates summary statistics of the polygons
    Parameters
    ----------
    gdf: GeoDataFrame

    Returns
    -------
    ss: GeoDataFrame
        summary statistics of the polygons
    cgdf: GeoDataFrame
        copy of the input GeoDataFrame with additional columns of the summary statistics
    """
    cgdf = gdf.copy()
    cgdf['primeter'] = cgdf.length
    cgdf['area'] = cgdf.area
    cgdf['ar_pratio'] = cgdf.area / cgdf.length
    # get the summary statics of the polygons
    ss = cgdf.quantile([0.25, 0.5, 0.75])
    return ss, cgdf

def buff_dfs(gdf: GeoDataFrame):
    """
    union and buffer the polygons of each class separately,
    to create continuous polygons and merge them into one GeoDataFrame.

    Parameters
    ----------
    gdf: GeoDataFrame
        Polygon dataframes with three classes in metric coordinate system
    crs: int

    Returns
    -------
    GeoDataFrame:
        merged GeoDataFrame of the three classes
    """

    gdf.geometry = gdf.simplify(0.2)
    dfsw = prepare_class_gdf(gdf, 'sidewalk')
    dfcw = prepare_class_gdf(gdf, 'crosswalk')
    dfrd = prepare_class_gdf(gdf, 'road')

    buffersw = buffer_union_erode(dfsw, 0.3, -0.3, 0.2, 0.3, 0.3)

    buffersw['f_type'] = 'sidewalk'

    buffercw = buffer_union_erode(dfcw, 0.3, -0.25, 0.2, 0.3, 0.3)

    buffercw['f_type'] = 'crosswalk'

    bufferrd = buffer_union_erode(dfrd, 0.4, -0.4, 0.2, 0.3, 0.3)
    bufferrd['f_type'] = 'road'

    merged = pd.concat([buffercw, buffersw, bufferrd])
    merged.geometry = merged.geometry.set_crs(gdf.crs)

    return merged


from regexParser import *
from styles import *
from docx import Document

# ALIGN CENTRE!

with open('md_test.md', 'r') as f:
    lines = f.read().strip().split('\n')
    print(lines)
    md = MarkdownParser()
    for l in lines:
        md.parseLine(l)
    doc = md.doc
    # paragraph = doc.add_paragraph("List 1", "List Bullet 2")
    paras = doc.paragraphs
    for p in paras:
        print(p.text)
    doc.save('doc_test.docx')


#!/usr/bin/env python2
# -*- coding: UTF-8 -*-

import unittest

yq_1_yq = '%Y.%m.%d_%H.%M.%S'
yq_3_yq = '%Y-%m'
yq_2_yq = yq_3_yq + '-%d'
yq_5_yq = '%Y%m%d%H%M%S'
yq_6_yq = '%H:%M'
yq_7_yq = yq_6_yq + ':%S'
yq_4_yq = ''.join([
    yq_2_yq,
    ' ',
    yq_7_yq,
    ])

class TestTimeFormats(unittest.TestCase):
    def test_time_formats(self):
        '''
        TestTimeFormats:
        '''
        self.assertEqual(yq_1_yq, '%Y.%m.%d_%H.%M.%S')
        self.assertEqual(yq_3_yq, '%Y-%m')
        self.assertEqual(yq_2_yq, '%Y-%m-%d')
        self.assertEqual(yq_5_yq, '%Y%m%d%H%M%S')
        self.assertEqual(yq_6_yq, '%H:%M')
        self.assertEqual(yq_7_yq, '%H:%M:%S')
        self.assertEqual(yq_4_yq, '%Y-%m-%d %H:%M:%S')


# Objetivo: Crie um programa onde o usuário possa digitar sete valores numéricos e cadastre-os em uma
# lista única que mantenha separados os valores pares e impares. no final mostre os valores pares e impares em ordem crescente.
# Data: 04/11/22
# Author: Washington

print(f"{' Lista de pares e ímpares ':-^50}", "\n")

pares = []
impares = []
lista_geral = [pares, impares]

for x in range(1, 8):
    num = int(input(f"Entre com o {x}º número: "))
    if num % 2 == 0:
        pares.append(num)
    else:
        impares.append(num)

pares.sort()
impares.sort()

print("-" * 50)

print(f"A lista geral é: {lista_geral}")
print(f"A lista dos númeors pares é : {pares}")
print(f"A lista dos números ímpares é: {impares}")

print("\n", f"{' Fim do Programa ':-^50}")


"""
Creates an artificial genome from a dataset file.

This method of producing a sequence of desired targets is based on code provided
in the supplementary information of Bradford and Perrin (2019).
"""

import consts
import common
from utils.datasets import DATASETS

# The default length of lines in the fasta file.
LINE_LEN = 50
# A template for the header of the fasta file.
HEADER = ">chr%d\n"

def get_args():
    parser = common.default_parser(
        "Creates an artificial genome from a dataset file.")
    parser.add_argument(
        "-l", "--line_len", type=int, default=LINE_LEN,
        help="Length of lines in the output fasta file")
    parser.add_argument(
        "-t", "--target_len", type=int, default=consts.TARGET_LEN,
        help="Length of target sequences")
    return parser.parse_args()

def break_lines(genome, line_len):
    """Breaks the sequence into short lines.

    Args:
        genome: The genome sequence.
        line_len: The length of the lines.

    Returns:
        A list of strings, each of length line_len (except for the last which
        could be shorter), that concatenated together form the genome.
    """
    return [genome[i:i+line_len] for i in range(0, len(genome), line_len)]

def main():
    """Creates an artificial genome from a dataset file.

    For command line help, run with the '-h' flag.

    Writes:
        An artificial genome fasta file along with a matching .txt file with the
        entire sequence in one line.
    """
    args = get_args()
    dataset = DATASETS[args.dataset]
    dataset.set_work_dir(args.path)
    chrom = dataset.get_chr(args.chr)

    targets = dataset.get_targets(args.chr, args.target_len)
    guide_gap = 'N' * consts.NUM_N_BETWEEN_GUIDES
    # Starts the sequence with padding, since some tools struggle with targets
    # at the beginning of the sequence.
    genome = guide_gap
    for target in targets:
        genome += target
        genome += guide_gap

    lines = break_lines(genome, args.line_len)
    # Writes the fasta file.
    with open(chrom.get_path(), 'w') as fd_fasta:
        fd_fasta.write(HEADER % args.chr)
        [fd_fasta.write(line + '\n') for line in lines]
    # Writes a txt file, with the entire sequence in one line.
    # The genome could be too long for writing in a single chunk, hence we use
    # the broken representation.
    with open(chrom.get_path('txt'), 'w') as fd_txt:
        [fd_txt.write(line) for line in lines]

if __name__ == '__main__':
    main()


import datetime
import json
import logging

import boto3
import custom_connector_sdk.connector.auth as auth
import custom_connector_sdk.connector.configuration as config
import custom_connector_sdk.connector.context as context
import custom_connector_sdk.connector.fields as fields
import custom_connector_sdk.connector.settings as settings
import custom_connector_sdk.lambda_handler.requests as requests
import custom_connector_sdk.lambda_handler.responses as responses
from custom_connector_sdk.lambda_handler.handlers import (
    ConfigurationHandler,
    MetadataHandler,
    RecordHandler,
)
from custom_connector_sdk.lambda_handler.lambda_handler import (
    BaseLambdaConnectorHandler,
)
from custom_connector_sdk.lambda_handler.responses import ErrorCode, ErrorDetails
from google.analytics.data_v1beta import BetaAnalyticsDataClient
from google.analytics.data_v1beta.types import (
    Dimension,
    Metric,
    OrderBy,
    RunRealtimeReportRequest,
)
from google.oauth2 import service_account

LOGGER = logging.getLogger()
LOGGER.setLevel(logging.INFO)

# カスタムコネクタの情報
CONNECTOR_OWNER = "DemoOwner"
CONNECTOR_NAME = "GA4Connector"
CONNECTOR_VERSION = "1.0"

# GA4 への接続情報のキー名（変更不要）
PROPERTY_ID_KEY = "propertyId"
PRIVATE_KEY_KEY = "privateKey"
CLIENT_EMAIL_KEY = "clientEmail"

# スキーマ定義（API 毎に定義）
# 今回は RunRealtimeReport API のみサポートするのでそのエンティティを定義
run_realtime_report_entity = context.Entity(
    entity_identifier="RunRealtimeReportEntity",
    label="RunRealtimeReportEntity",
    has_nested_entities=False,
    description="RunRealtimeReportEntity",
)

class ConnectorDemoConfigurationHandler(ConfigurationHandler):
    def validate_connector_runtime_settings(
        self, request: requests.ValidateConnectorRuntimeSettingsRequest
    ) -> responses.ValidateConnectorRuntimeSettingsResponse:
        print("ConfigurationHandler::validate_connector_runtime_settings")
        return responses.ValidateConnectorRuntimeSettingsResponse(is_success=True)

    def validate_credentials(
        self, request: requests.ValidateCredentialsRequest
    ) -> responses.ValidateCredentialsResponse:
        print("ConfigurationHandler::validate_credentials")
        return responses.ValidateCredentialsResponse(is_success=True)

    # このメソッドが CustomConnector 登録時に呼び出される
    # 実装内容は連携先システムへの認証方法の定義
    def describe_connector_configuration(
        self, request: requests.DescribeConnectorConfigurationRequest
    ) -> responses.DescribeConnectorConfigurationResponse:
        print("ConfigurationHandler::describe_connector_configuration")

        # コネクタの認証方法の設定
        # 今回は GA4 Data API を叩く際に GCP のサービスアカウントを使用するため Custom Auth を設定
        authentication_config = auth.AuthenticationConfig(
            is_custom_auth_supported=True,
            custom_auth_config=[
                auth.CustomAuthConfig(
                    authentication_type="CUSTOM",
                    auth_parameters=[
                        auth.AuthParameter(
                            key=PRIVATE_KEY_KEY,
                            required=True,
                            label="Private Key",
                            description="GCP サービスアカウントの private_key を入力（スペースは \\s に置換）",
                            sensitive_field=True,
                            connector_supplied_values=None,
                        ),
                        auth.AuthParameter(
                            key=CLIENT_EMAIL_KEY,
                            required=True,
                            label="Client Email",
                            description="xxxxxx@xxxx.iam.gserviceaccount.com",
                            sensitive_field=False,
                            connector_supplied_values=None,
                        ),
                    ],
                )
            ],
        )

        property_id_setting = settings.ConnectorRuntimeSetting(
            key=PROPERTY_ID_KEY,
            data_type=settings.ConnectorRuntimeSettingDataType.String,
            required=True,
            label="GA4 Property ID",
            description="GA4 Property ID",
            scope=settings.ConnectorRuntimeSettingScope.CONNECTOR_PROFILE,
        )

        return responses.DescribeConnectorConfigurationResponse(
            is_success=True,
            # カスタムコネクタの所有者
            connector_owner=CONNECTOR_OWNER,
            # カスタムコネクタ名
            connector_name=CONNECTOR_NAME,
            # 認証設定
            authentication_config=authentication_config,
            # このカスタムコネクタのバージョン
            connector_version=CONNECTOR_VERSION,
            # サポートしている GA4 Data API のバージョン
            supported_api_versions=["GoogleAnalyticsDataAPIv1"],
            # カスタムコネクタのモードを設定
            # このカスタムコネクタは送信元にのみ指定できるため, ConnectorModes.SOURCE のみ設定
            connector_modes=[config.ConnectorModes.SOURCE],
            # コネクタの runtime の設定
            # 今回は GA4 のプロパティ ID を設定
            connector_runtime_setting=[property_id_setting],
        )

class ConnectorDemoRecordHandler(RecordHandler):
    def retrieve_data(
        self, request: requests.RetrieveDataRequest
    ) -> responses.RetrieveDataResponse:
        print("RecordHandler::retrieve_data")
        record_list = []
        return responses.RetrieveDataResponse(is_success=True, records=record_list)

    def write_data(
        self, request: requests.WriteDataRequest
    ) -> responses.WriteDataResponse:
        print("RecordHandler::write_data")
        write_record_results = []

        return responses.WriteDataResponse(
            is_success=True, write_record_results=write_record_results
        )

    def query_data(
        self, request: requests.QueryDataRequest
    ) -> responses.QueryDataResponse:
        print("RecordHandler::query_data")

        # コネクタの runtime に GA4 のプロパティ ID が設定されているか検証
        if PROPERTY_ID_KEY not in request.connector_context.connector_runtime_settings:
            error_message = f"{PROPERTY_ID_KEY} should be provided as runtime setting"
            LOGGER.error(f"QueryData request failed with entity: {error_message}")
            return responses.QueryDataResponse(
                is_success=False,
                error_details=ErrorDetails(
                    error_code=ErrorCode.InvalidArgument,
                    error_message=error_message,
                ),
            )

        # 有効なエンティティか検証（ここでは RunRealtimeReportEntity かどうか検証）
        entity_id = request.connector_context.entity_definition.entity.entity_identifier
        if entity_id != run_realtime_report_entity.entity_identifier:
            error_message = f"{entity_id} is not valid entity"
            LOGGER.error(f"QueryData request failed with entity: {error_message}")
            return responses.QueryDataResponse(
                is_success=False,
                error_details=ErrorDetails(
                    error_code=ErrorCode.InvalidArgument,
                    error_message=error_message,
                ),
            )

        # AppFlow の接続作成時に設定した認証情報を Secrets Manager から取得
        secrets = boto3.client("secretsmanager").get_secret_value(
            SecretId=request.connector_context.credentials.secret_arn
        )
        secret_string = json.loads(secrets["SecretString"])

        credentials = service_account.Credentials.from_service_account_info(
            {
                "private_key": secret_string[PRIVATE_KEY_KEY]
                .replace("\\s", " ")
                .replace("\\n", "\n"),
                "client_email": secret_string[CLIENT_EMAIL_KEY],
                "token_uri": "https://oauth2.googleapis.com/token",
            }
        )

        # GA4 Data API を叩く
        client = BetaAnalyticsDataClient(credentials=credentials)

        req = RunRealtimeReportRequest(
            property=f"properties/{request.connector_context.connector_runtime_settings[PROPERTY_ID_KEY]}",
            # 分析項目の指定
            dimensions=[
                Dimension(name="minutesAgo"),
                Dimension(name="country"),
                Dimension(name="city"),
            ],
            # 集計項目の指定
            metrics=[Metric(name="screenPageViews"), Metric(name="activeUsers")],
            # 順序の指定
            order_bys=[
                OrderBy(
                    dimension=OrderBy.DimensionOrderBy(dimension_name="minutesAgo"),
                    desc=False,
                )
            ],
            # 取得する行数 (最大100,000)
            limit=100000,
        )

        res = client.run_realtime_report(req)

        # 秒を切り捨てた現在時刻を取得
        now_floor_str = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:00")
        now = datetime.datetime.strptime(now_floor_str, "%Y-%m-%d %H:%M:%S")

        # JSON オブジェクトの文字列の配列
        # RunRealtimeReport のエンティティでマッピングしたデータフィールドだけ送信される
        # record_list = [
        #     '{"DateTime": "2022-01-01 00:00:00", "RequestDateTime": "2022-01-01 00:00:00", "ScreenPageViews": 1}',
        #     '{"DateTime": "2022-01-01 00:01:00", "RequestDateTime": "2022-01-01 00:01:00", "ScreenPageViews": 2}',
        # ]
        record_list = []

        for row in res.rows:
            # イベント発生日時の絶対値をイベント発生日時の相対値 (dimensions.minutesAgo) から割り出す
            time = now - datetime.timedelta(minutes=int(row.dimension_values[0].value))

            record_list.append(
                json.dumps(
                    {
                        # レポートをリクエストした時間
                        # * DIMENSION が同じで METRIC が異なるデータは RequestDateTime が新しい方が正確なデータ
                        "RequestDateTime": now.strftime("%Y-%m-%d %H:%M:00"),
                        # [DIMENSION]: イベントの発生日時
                        "DateTime": time.strftime("%Y-%m-%d %H:%M:00"),
                        # [DIMENSION]: イベントを発生させた訪問者の国名
                        "Country": row.dimension_values[1].value,
                        # [DIMENSION]: イベントを発生させた訪問者の町名
                        "City": row.dimension_values[2].value,
                        # [METRIC]: dimensions に対応する訪問者がページを開いた回数
                        "ScreenPageViews": row.metric_values[0].value,
                        # [METRIC]: dimensions に対応するアクティブな訪問者の数
                        "ActiveUsers": row.metric_values[1].value,
                    }
                )
            )

        return responses.QueryDataResponse(is_success=True, records=record_list)

class ConnectorDemoMetadataHandler(MetadataHandler):
    # フロー作成時に呼び出され、スキーマ定義のリストを返す
    def list_entities(
        self, request: requests.ListEntitiesRequest
    ) -> responses.ListEntitiesResponse:
        print("MetadataHandler::list_entities")

        # フロー作成時のスキーマ定義
        # 今回は RunRealtimeReport API のみサポートしているため 1 つのエンティティを設定
        entity_list = [run_realtime_report_entity]

        return responses.ListEntitiesResponse(is_success=True, entities=entity_list)

    # フロー作成のデータフィールドマッピングで呼び出され、フィールドの定義を返す
    def describe_entity(
        self, request: requests.DescribeEntityRequest
    ) -> responses.DescribeEntityResponse:
        print("MetadataHandler::describe_entity")

        date_time_field = context.FieldDefinition(
            field_name="DateTime",
            data_type=fields.FieldDataType.String,
            data_type_label="string",
            label="DateTime",
            description="DateTime",
            default_value="1970-01-01 00:00:00",
            is_primary_key=True,
            # データの送信元に指定する場合に設定するフィールドの型定義
            read_properties=fields.ReadOperationProperty(
                # データ取得時にフィルタとして指定できる項目かどうか
                is_queryable=True,
                # データ取得時に表示されるかどうか
                is_retrievable=True,
                # null 値を許容するかどうか
                is_nullable=False,
                # Date か DateTime 型かどうか
                is_timestamp_field_for_incremental_queries=False,
            ),
            # データの送信先に指定する場合に設定するフィールドの型定義
            write_properties=None,
        )

        request_date_time_field = context.FieldDefinition(
            field_name="RequestDateTime",
            data_type=fields.FieldDataType.String,
            data_type_label="string",
            label="RequestDateTime",
            description="RequestDateTime",
            default_value="1970-01-01 00:00:00",
            is_primary_key=False,
            read_properties=fields.ReadOperationProperty(
                is_queryable=True,
                is_retrievable=True,
                is_nullable=False,
                is_timestamp_field_for_incremental_queries=False,
            ),
            write_properties=None,
        )

        country_field = context.FieldDefinition(
            field_name="Country",
            data_type=fields.FieldDataType.String,
            data_type_label="string",
            label="Country",
            description="Country",
            default_value="Country",
            is_primary_key=False,
            read_properties=fields.ReadOperationProperty(
                is_queryable=True,
                is_retrievable=True,
                is_nullable=False,
                is_timestamp_field_for_incremental_queries=False,
            ),
            write_properties=None,
        )

        city_field = context.FieldDefinition(
            field_name="City",
            data_type=fields.FieldDataType.String,
            data_type_label="string",
            label="City",
            description="City",
            default_value="City",
            is_primary_key=False,
            read_properties=fields.ReadOperationProperty(
                is_queryable=True,
                is_retrievable=True,
                is_nullable=False,
                is_timestamp_field_for_incremental_queries=False,
            ),
            write_properties=None,
        )

        screen_page_views_field = context.FieldDefinition(
            field_name="ScreenPageViews",
            data_type=fields.FieldDataType.Integer,
            data_type_label="int",
            label="ScreenPageViews",
            description="ScreenPageViews",
            default_value="0",
            is_primary_key=False,
            read_properties=fields.ReadOperationProperty(
                is_queryable=True,
                is_retrievable=True,
                is_nullable=False,
                is_timestamp_field_for_incremental_queries=False,
            ),
            write_properties=None,
        )

        active_users_field = context.FieldDefinition(
            field_name="ActiveUsers",
            data_type=fields.FieldDataType.Integer,
            data_type_label="int",
            label="ActiveUsers",
            description="ActiveUsers",
            default_value="0",
            is_primary_key=False,
            read_properties=fields.ReadOperationProperty(
                is_queryable=True,
                is_retrievable=True,
                is_nullable=False,
                is_timestamp_field_for_incremental_queries=False,
            ),
            write_properties=None,
        )

        # 今回は RunRealtimeReport API の DataTime, RequestDatetime, ScreenPageViews を fields に設定
        entity_definition = context.EntityDefinition(
            entity=run_realtime_report_entity,
            fields=[
                date_time_field,
                request_date_time_field,
                country_field,
                city_field,
                screen_page_views_field,
                active_users_field,
            ],
        )

        return responses.DescribeEntityResponse(
            is_success=True, entity_definition=entity_definition
        )

class ConnectorDemoLambdaHandler(BaseLambdaConnectorHandler):
    def __init__(self):
        super().__init__(
            ConnectorDemoMetadataHandler(),
            ConnectorDemoRecordHandler(),
            ConnectorDemoConfigurationHandler(),
        )

def handler(event, context):
    return ConnectorDemoLambdaHandler().lambda_handler(event, context)