# Copyright (c) 2020 Matthew Earl
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
#     The above copyright notice and this permission notice shall be included
#     in all copies or substantial portions of the Software.
#
#     THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
#     OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
#     MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN
#     NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM,
#     DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR
#     OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE
#     USE OR OTHER DEALINGS IN THE SOFTWARE.

import collections
import dataclasses
import functools
import logging
import math
import re
import time
from dataclasses import dataclass
from typing import List, Tuple, Dict

import bpy
import bpy_types
import numpy as np

from . import proto, bsp, mdl, blendmdl, blendbsp, simplex, blendpart

logger = logging.getLogger(__name__)
Vec3 = Tuple[float, float, float]

_NEAR_CLIP_PLANE = 8
_FAR_CLIP_PLANE = 2048
_EYE_HEIGHT = 22

def _patch_vec(old_vec: Vec3, update):
    return tuple(v if u is None else u for v, u in zip(old_vec, update))

def _quake_to_blender_angles(quake_angles: Vec3) -> Vec3:
    return (math.pi * (90. - quake_angles[0]) / 180,
            math.pi * quake_angles[2] / 180,
            math.pi * (quake_angles[1] - 90.) / 180)

def _fix_angles(old_angles, new_angles, degrees=False):
    # Undo wrapping of the yaw
    if degrees:
        t = 360
    else:
        t = 2 * np.pi

    old_yaw = old_angles[1] / t
    new_yaw = new_angles[1] / t
    i = old_yaw // 1
    j = np.argmin(np.abs(np.array([i - 1, i, i + 1]) + new_yaw - old_yaw))
    return (new_angles[0], (new_yaw + i + j - 1) * t, new_angles[2])

def _quake_angles_to_mat(angles):
    # Ported from mathlib.c AngleVectors
    pitch = angles[0] * (np.pi / 180);
    yaw = angles[1] * (np.pi / 180);
    roll = angles[2] * (np.pi / 180);

    sy, cy = np.sin(yaw), np.cos(yaw)
    sp, cp = np.sin(pitch), np.cos(pitch)
    sr, cr = np.sin(roll), np.cos(roll)

    right = np.array([-1*sr*sp*cy + -1*cr*-sy, -1*sr*sp*sy + -1*cr*cy, -1*sr*cp])
    forward = np.array([cp*cy, cp*sy, -sp])
    up = np.array([cr*sp*cy + -sr*-sy, cr*sp*sy + -sr*cy, cr*cp])

    return np.stack([right, forward, up], axis=1)

def _get_model_config(mdl_name, config):
    mdls_cfg = config['models']
    cfg = dict(mdls_cfg['__default__'])
    cfg.update(mdls_cfg.get(mdl_name, {}))
    return cfg

@dataclass
class _EntityInfo:
    """Information for a single entity slot"""
    model_num: int
    frame: int
    skin: int
    origin: Vec3
    angles: Vec3

    def update(self, msg: proto.ServerMessageUpdate, baseline: "_EntityInfo"):
        def none_or(a, b):
            return a if a is not None else b

        angles = _fix_angles(self.angles, _patch_vec(baseline.angles, msg.angle))

        return dataclasses.replace(
            baseline,
            model_num=none_or(msg.model_num, baseline.model_num),
            frame=none_or(msg.frame, baseline.frame),
            skin=none_or(msg.skin, baseline.skin),
            origin=_patch_vec(baseline.origin, msg.origin),
            angles=angles
        )

_DEFAULT_BASELINE = _EntityInfo(0, 0, 0, (0, 0, 0), (0, 0, 0))

class SampleAsLightObject:
    @property
    def bbox(self):
        raise NotImplementedError

    @property
    def leaf(self):
        raise NotImplementedError

    @property
    def origin(self):
        raise NotImplementedError

    def add_keyframe(self, vis: bool, blender_frame: int):
        raise NotImplementedError

class AliasModelSampleAsLightObject:
    _bm: blendmdl.BlendMdl
    _bb: blendbsp.BlendBsp

    def __init__(self, bm, bb, mdl_cfg):
        self._bm = bm
        self._bb = bb
        self._bbox = np.array(mdl_cfg['bbox'])

    @property
    def bbox(self):
        return np.array(self._bm.obj.location) + self._bbox

    @property
    def leaf(self):
        return self._bb.bsp.models[0].get_leaf_from_point(self._bm.obj.location)

    def add_keyframe(self, vis: bool, blender_frame: int):
        for bmat in self._bm.sample_as_light_mats:
            bmat.add_sample_as_light_keyframe(vis, blender_frame)

@dataclass(eq=False)
class LeafSampleAsLightObject:
    _leaf: bsp.Leaf
    _mat: bpy.types.Material
    _tex_cfg: Dict
    _model_idx: int
    _bb: blendbsp.BlendBsp

    @property
    def bbox(self):
        if "bbox" not in self._tex_cfg:
            raise Exception("Sample as light textures must have bounding boxes")
        tex_bbox = np.array(self._tex_cfg['bbox'])
        return np.stack([self._leaf.bbox.mins + tex_bbox[0],
                         self._leaf.bbox.maxs + tex_bbox[1]]) + self._model_origin

    @property
    def _model(self):
        return self._bb.bsp.models[self._model_idx]

    @property
    def _model_origin(self):
        return np.array(self._bb.model_objs[self._model].location)

    @property
    def leaf(self):
        if self._model_idx == 0:
            out = self._leaf
        else:
            leaf_origin = 0.5 * (np.array(self._leaf.bbox.mins) +
                                 np.array(self._leaf.bbox.maxs))
            origin = leaf_origin + self._model_origin
            out = self._bb.bsp.models[0].get_leaf_from_point(origin)
        return out

    def add_keyframe(self, vis: bool, blender_frame: int):
        self._mat.add_sample_as_light_keyframe(vis, blender_frame)

    @classmethod
    def create_from_bsp(cls, bb: blendbsp.BlendBsp):
        if bb.sample_as_light_info:
            return (
                cls(leaf, mat, tex_cfg, model.id_, bb)
                for model, model_info in bb.sample_as_light_info.items()
                for leaf, leaf_info in model_info.items()
                for mat, tex_cfg in leaf_info.items()
            )
        else:
            return []

@dataclass
class ManagedObject:
    fps: float

    def _get_blender_frame(self, time):
        return int(round(self.fps * time))

    @property
    def ignore_duplicate_updates(self) -> bool:
        raise NotImplementedError

    def add_pose_keyframe(self, pose_num: int, time: float):
        raise NotImplementedError

    def add_visible_keyframe(self, visible: bool, time: float):
        raise NotImplementedError

    def add_origin_keyframe(self, origin: Vec3, time: float):
        raise NotImplementedError

    def add_angles_keyframe(self, angles: Vec3, time: float):
        raise NotImplementedError

    def set_invisible_to_camera(self):
        raise NotImplementedError

    def done(self, final_time: float):
        raise NotImplementedError

@dataclass
class AliasModelManagedObject(ManagedObject):
    bm: blendmdl.BlendMdl

    @property
    def ignore_duplicate_updates(self) -> bool:
        return True

    def add_pose_keyframe(self, pose_num: int, time: float):
        self.bm.add_pose_keyframe(pose_num, time, self.fps)

    def add_visible_keyframe(self, visible: bool, time: float):
        blender_frame = self._get_blender_frame(time)
        for sub_obj in self.bm.sub_objs:
            sub_obj.hide_render = not visible
            sub_obj.keyframe_insert('hide_render', frame=blender_frame)
            sub_obj.hide_viewport = not visible
            sub_obj.keyframe_insert('hide_viewport', frame=blender_frame)

    def add_origin_keyframe(self, origin: Vec3, time: float):
        self.bm.obj.location = origin
        self.bm.obj.keyframe_insert('location', frame=self._get_blender_frame(time))

    def add_angles_keyframe(self, angles: Vec3, time: float):
        # Should this use the other angles?
        self.bm.obj.rotation_euler = (0., 0., angles[1])
        if self.bm.am.header['flags'] & mdl.ModelFlags.ROTATE:
            self.bm.obj.rotation_euler.z = time * 100. * np.pi / 180
        self.bm.obj.keyframe_insert('rotation_euler', frame=self._get_blender_frame(time))

    def set_invisible_to_camera(self):
        self.bm.set_invisible_to_camera()

    def done(self, final_time: float):
        self.bm.done(final_time, self.fps)

@dataclass
class BspModelManagedObject(ManagedObject):
    _bb: blendbsp.BlendBsp
    _model_num: int

    @property
    def ignore_duplicate_updates(self) -> bool:
        return False

    def add_pose_keyframe(self, pose_num: int, time: float):
        model = self._bb.bsp.models[self._model_num]
        self._bb.add_material_frame_keyframe(model, pose_num, self._get_blender_frame(time))

    @property
    def _model(self):
        return self._bb.bsp.models[self._model_num]

    def add_visible_keyframe(self, visible: bool, time: float):
        blender_frame = self._get_blender_frame(time)
        self._bb.add_visible_keyframe(self._model, visible, blender_frame)

    def add_origin_keyframe(self, origin: Vec3, time: float):
        obj = self._bb.model_objs[self._model]
        obj.location = origin
        obj.keyframe_insert('location', frame=self._get_blender_frame(time))

    def add_angles_keyframe(self, angles: Vec3, time: float):
        pass

    def done(self, final_time: float):
        self._bb.add_animated_material_keyframes(self._get_blender_frame(final_time), final_time)

@dataclass
class NullManagedObject(ManagedObject):
    @property
    def ignore_duplicate_updates(self) -> bool:
        return False

    def add_pose_keyframe(self, pose_num: int, time: float):
        pass

    def add_visible_keyframe(self, visible: bool, time: float):
        pass

    def add_origin_keyframe(self, origin: Vec3, time: float):
        pass

    def add_angles_keyframe(self, angles: Vec3, time: float):
        pass

    def done(self, final_time: float):
        pass

class ObjectManager:
    def __init__(self, fs, config, fps, fov, width, height, world_obj_name='demo', load_level=True):
        self._fs = fs
        self._fps = fps
        self._config = config
        self._load_level = load_level

        self._pal = np.fromstring(fs['gfx/palette.lmp'], dtype=np.uint8).reshape(256, 3) / 255

        self._bb: Optional[blendbsp.BlendBsp] = None
        self._objs: Dict[Tuple[int, int], ManagedObject] = {}
        self._model_paths: Optional[List[str]] = None
        self._view_entity_num: Optional[int] = None
        self._static_objects: List[ManagedObject] = []
        self._sample_as_light_objects: List[SampleAsLightObject] = []
        self._sal_time: float = 0.
        self._first_update_time: Optional[float] = None
        self._intermission = False
        self._num_explosions: int = 0
        self._num_teleports: int = 0

        self.world_obj = bpy.data.objects.new(world_obj_name, None)
        bpy.context.scene.collection.objects.link(self.world_obj)

        blendpart.get_particle_root().parent = self.world_obj

        self._width, self._height = width, height
        bpy.data.scenes['Scene'].render.resolution_x = width
        bpy.data.scenes['Scene'].render.resolution_y = height

        self._fov = fov
        demo_cam = bpy.data.cameras.new(name="demo_cam")
        demo_cam.angle = fov * np.pi / 180
        demo_cam.clip_start = 0.04
        self._demo_cam_obj = bpy.data.objects.new(name="demo_cam", object_data=demo_cam)
        bpy.context.scene.collection.objects.link(self._demo_cam_obj)
        self._demo_cam_obj.parent = self.world_obj

    def set_intermission(self, i: bool):
        self._intermission = i

    def _path_to_bsp_name(self, bsp_path):
        m = re.match(r"maps/([a-zA-Z0-9_]+).bsp", bsp_path)
        if m is None:
            raise Exception("Unexpected BSP path {mdl_path}")
        return m.group(1)

    def set_model_paths(self, model_paths: List[str]):
        if self._model_paths is not None:
            raise Exception("Model paths already set")
        self._model_paths = model_paths

        if self._load_level:
            map_path = self._model_paths[0]
            logger.info('Parsing bsp %s', map_path)
            b = bsp.Bsp(self._fs.open(map_path))
            map_name = self._path_to_bsp_name(map_path)
            logger.info('Adding bsp %s', map_path)
            self._bb = blendbsp.add_bsp(b, self._pal, map_name, self._config)
            self._bb.map_obj.parent = self.world_obj
            self._bb.hide_all_but_main()

            self._sample_as_light_objects.extend(
                LeafSampleAsLightObject.create_from_bsp(self._bb)
            )

    def _path_to_model_name(self, mdl_path):
        m = re.match(r"progs/([A-Za-z0-9-_]*)\.mdl", mdl_path)
        if m is None:
            raise Exception("Unexpected model path {mdl_path}")
        return m.group(1)

    @functools.lru_cache(None)
    def _load_alias_model(self, model_path):
        return mdl.AliasModel(self._fs.open(model_path))

    def set_view_entity(self, entity_num):
        self._view_entity_num = entity_num

    def create_static_object(self, model_num, frame, origin, angles, skin):
        model_path = self._model_paths[model_num - 1]
        am = mdl.AliasModel(self._fs.open(model_path))

        mdl_name = self._path_to_model_name(model_path)
        mdl_cfg = _get_model_config(mdl_name, self._config)
        bm = blendmdl.add_model(am,
                                self._pal,
                                mdl_name,
                                f"static{len(self._static_objects)}",
                                skin,
                                mdl_cfg,
                                frame,
                                self._config['do_materials'])
        bm.obj.parent = self.world_obj
        bm.obj.location = origin
        bm.obj.rotation_euler = (0., 0., angles[1])

        self._static_objects.append(bm)

        if bm.sample_as_light_mats:
            self._sample_as_light_objects.append(
                AliasModelSampleAsLightObject(bm, self._bb, mdl_cfg)
            )

    def create_teleport(self, pos, time):
        obj_name = f'teleport{self._num_teleports}'
        blendpart.create_teleport(time, obj_name, pos, self._fps)
        self._num_teleports += 1

    def create_explosion(self, pos, time):
        obj_name = f'explosion{self._num_explosions}'
        blendpart.create_explosion(time, obj_name, pos, self._fps)
        self._num_explosions += 1

    def _create_managed_object(self, entity_num, model_num, skin_num, initial_pose_num):
        model_path = self._model_paths[model_num - 1] if model_num != 0 else None

        if model_num == 0:
            # Used to make objects disappear, eg. player at the end of the level
            managed_obj = NullManagedObject(self._fps)
        elif model_path.startswith('*'):
            if self._load_level:
                map_model_idx = int(model_path[1:])
                managed_obj = BspModelManagedObject(self._fps, self._bb, map_model_idx)
            else:
                managed_obj = NullManagedObject(self._fps)
        elif model_path.endswith('.mdl'):
            am = self._load_alias_model(model_path)
            mdl_name = self._path_to_model_name(model_path)
            logger.info('Loading alias model %s', mdl_name)
            mdl_cfg = _get_model_config(mdl_name, self._config)
            bm = blendmdl.add_model(am,
                                    self._pal,
                                    mdl_name,
                                    f'ent{entity_num}_{mdl_name}',
                                    skin_num,
                                    mdl_cfg,
                                    initial_pose_num,
                                    self._config['do_materials'])
            bm.obj.parent = self.world_obj
            managed_obj = AliasModelManagedObject(self._fps, bm)

            if bm.sample_as_light_mats:
                self._sample_as_light_objects.append(
                    AliasModelSampleAsLightObject(bm, self._bb, mdl_cfg)
                )

        elif model_path.endswith('.bsp'):
            bsp_name = self._path_to_bsp_name(model_path)
            logger.info('Loading bsp model %s', bsp_name)
            b = bsp.Bsp(self._fs.open(model_path))
            if len(b.models) != 1:
                raise Exception(f"Expected one model in bsp model {bsp_name}, not {len(b.models)}")
            bb = blendbsp.add_bsp(b, self._pal, bsp_name, self._config,
                                  f'ent{entity_num}_')
            bb.map_obj.parent = self.world_obj
            managed_obj = BspModelManagedObject(self._fps, bb, 0)
        else:
            logging.warning('Cannot handle model %r', model_path)
            managed_obj = NullManagedObject(self._fps)

        return managed_obj

    def _view_simplex(self, view_origin, view_angles):
        view_origin = np.array(view_origin)

        aspect_ratio = self._width / self._height
        tan_fov = np.tan(0.5 * self._fov * np.pi / 180)
        if aspect_ratio > 1:
            h_tan = tan_fov
            v_tan = h_tan / aspect_ratio
        else:
            v_tan = tan_fov
            h_tan = v_tan * aspect_ratio

        constraints = np.array([
            [-1, h_tan, 0, 0],               # right
            [1, h_tan, 0, 0],                # left
            [0, v_tan, 1, 0],                # bottom
            [0, v_tan, -1, 0],               # top
            [0, 1, 0, -_NEAR_CLIP_PLANE],    # near
            [0, -1, 0, _FAR_CLIP_PLANE],     # far
        ])

        constraints[:, :3] /= np.linalg.norm(constraints[:, :3], axis=1)[:, None]
        rotation_matrix = _quake_angles_to_mat(view_angles)
        constraints[:, :3] = constraints[:, :3] @ rotation_matrix.T
        constraints[:, 3] -= np.einsum('ij,j->i', constraints[:, :3], view_origin)

        return simplex.Simplex(3, constraints, np.array([False, True, False, True, False, True]))

    def _simplex_bbox_test(self, bbox: np.ndarray, sx: simplex.Simplex):
        bbox_simplex = simplex.Simplex.from_bbox(*bbox)
        try:
            bbox_simplex.intersect(sx)
        except simplex.Infeasible:
            intersect = False
        else:
            intersect = True
        return intersect

    def _update_sample_as_light(self, view_origin, view_angles, blender_frame, crude_test=True):
        if not self._load_level:
            # No PVS info so can't update sample as lights
            return

        start = time.perf_counter()
        view_pvs = set(self._bb.bsp.models[0].get_leaf_from_point(view_origin).visible_leaves)

        view_sx = self._view_simplex(view_origin, view_angles)

        num_tests = 0
        num_visible = 0
        num_early_exits = 0
        for sal_obj in self._sample_as_light_objects:
            # Checking if the light PVS intersects with the view PVS.
            pvs = view_pvs & set(sal_obj.leaf.visible_leaves)
            vis = bool(pvs)

            # Clip leaf PVS bboxes to the light bbox.
            if vis:
                leaf_bboxes = np.stack([[leaf.bbox.mins, leaf.bbox.maxs] for leaf in pvs])
                bboxes = np.stack([
                    np.maximum(leaf_bboxes[:, 0, :], sal_obj.bbox[0][None, :]),
                    np.minimum(leaf_bboxes[:, 1, :], sal_obj.bbox[1][None, :])
                ], axis=1)
                bboxes = bboxes[np.all(bboxes[:, 0] < bboxes[:, 1], axis=1)]
                vis = bboxes.shape[0] != 0

            # Test if a single bbox that bounds all of the above bboxes intersects the view
            # frustum.
            if vis:
                crude_bbox = np.stack([
                    np.min(bboxes[:, 0, :], axis=0),
                    np.max(bboxes[:, 1, :], axis=0),
                ])
                num_tests += 1
                vis = self._simplex_bbox_test(crude_bbox, view_sx)

            # Finally, check if any of the individual bboxes intersects the view frustum.
            if not crude_test and vis:
                for bbox in bboxes:
                    num_tests += 1
                    if self._simplex_bbox_test(bbox, view_sx):
                        break
                else:
                    vis = False
            else:
                num_early_exits += 1

            num_visible += vis
            sal_obj.add_keyframe(vis, blender_frame)

        self._sal_time += time.perf_counter() - start
        logger.debug('frame: %s, frustum tests: %s, lights visible: %s, early exits: %s / %s, total sal time: %s',
                     blender_frame, num_tests, num_visible,
                     num_early_exits, len(self._sample_as_light_objects),
                     self._sal_time)

    def update(self, time, prev_entities, entities, prev_updated, updated, view_angles):
        blender_frame = int(round(self._fps * time))

        if self._intermission:
            view_angles = tuple(x * 180 / np.pi for x in entities[self._view_entity_num].angles)

        # Hide any objects that weren't updated in this frame, or whose model changed.
        for entity_num in prev_updated:
            obj = self._objs[entity_num, prev_entities[entity_num].model_num]
            if (prev_entities[entity_num].model_num != entities[entity_num].model_num
                    or entity_num not in updated):
                obj.add_visible_keyframe(False, time)

                # Insert keyframes so the object doesn't drift towards where the entity is re-used.
                obj.add_origin_keyframe(prev_entities[entity_num].origin, time)
                obj.add_angles_keyframe(prev_entities[entity_num].angles, time)

        for entity_num in updated:
            # Create managed objects where we don't already have one for the given entity num / model num.
            ent = entities[entity_num]
            model_num = entities[entity_num].model_num
            key = entity_num, model_num
            if key not in self._objs:
                obj = self._create_managed_object(entity_num, model_num, ent.skin, ent.frame)
                obj.add_visible_keyframe(False, 0)
                self._objs[key] = obj
            else:
                obj = self._objs[key]

            # Update position / rotation / pose
            prev_ent = prev_entities.get(entity_num)
            if (not obj.ignore_duplicate_updates or
                    prev_ent is None or
                    prev_ent.origin != ent.origin or
                    prev_ent.angles != ent.angles):
                obj.add_origin_keyframe(ent.origin, time)
                obj.add_angles_keyframe(ent.angles, time)
            obj.add_pose_keyframe(ent.frame, time)

        # Unhide objects that were updated this frame, or whose model changed.
        for entity_num in updated:
            if (entity_num not in prev_updated or
                (entity_num in prev_entities and
                 prev_entities[entity_num].model_num != entities[entity_num].model_num)):
                self._objs[entity_num, entities[entity_num].model_num].add_visible_keyframe(
                    True, time
                )

        # Pose camera
        view_origin = entities[self._view_entity_num].origin
        if not self._intermission:
            view_origin = (view_origin[0], view_origin[1], view_origin[2] + _EYE_HEIGHT)
        self._demo_cam_obj.location = view_origin
        self._demo_cam_obj.keyframe_insert('location', frame=blender_frame)
        self._demo_cam_obj.rotation_euler = _quake_to_blender_angles(view_angles)
        self._demo_cam_obj.keyframe_insert('rotation_euler', frame=blender_frame)

        # Set sample_as_light materials.
        self._update_sample_as_light(view_origin, view_angles, blender_frame)

        # Record for updating latest
        if self._first_update_time is None:
            self._first_update_time = time

    def done(self, final_time: float):
        # Animate static objects
        for bm in self._static_objects:
            bm.done(final_time, self._fps)

        for obj in self._objs.values():
            obj.done(final_time)

        # Make the view entity invisible to camera rays
        if self._config['hide_view_entity']:
            for (entity_num, model_num), obj in self._objs.items():
                if entity_num == self._view_entity_num and model_num != 0:
                    obj.set_invisible_to_camera()

        # Set start / end frame
        if self._first_update_time is not None:
            bpy.data.scenes['Scene'].frame_start = int(round(self._first_update_time * self._fps))
        bpy.data.scenes['Scene'].frame_end = int(round(final_time * self._fps))

def add_demo(demo_file, fs, config, fps=30, world_obj_name='demo',
             load_level=True, relative_time=False, fov=120, width=1920, height=1080):
    assert not relative_time, "Not yet supported"

    baseline_entities: Dict[int, _EntityInfo] = collections.defaultdict(lambda: _DEFAULT_BASELINE)
    entities: Dict[int, _EntityInfo] = {}
    fixed_view_angles: Vec3 = (0, 0, 0)
    prev_updated = set()
    demo_done = False
    obj_mgr = ObjectManager(fs, config, fps, fov, width, height, world_obj_name, load_level)
    last_time = 0.

    msg_iter = proto.read_demo_file(demo_file)

    while not demo_done:
        time = None
        update_done = False
        updated = set()
        prev_entities = dict(entities)

        while not update_done and not demo_done:
            try:
                msg_end, view_angles, parsed = next(msg_iter)
            except StopIteration:
                demo_done = True
                break

            update_done = msg_end and time is not None and entities

            fixed_view_angles = _fix_angles(fixed_view_angles, view_angles, degrees=True)

            if parsed.msg_type == proto.ServerMessageType.TIME:
                time = parsed.time

            if parsed.msg_type == proto.ServerMessageType.SERVERINFO:
                obj_mgr.set_model_paths(parsed.models)

            if parsed.msg_type == proto.ServerMessageType.SETVIEW:
                obj_mgr.set_view_entity(parsed.viewentity)

            if parsed.msg_type == proto.ServerMessageType.SETANGLE:
                fixed_view_angles = parsed.view_angles

            if parsed.msg_type == proto.ServerMessageType.SPAWNSTATIC:
                obj_mgr.create_static_object(
                    parsed.model_num, parsed.frame, parsed.origin, parsed.angles, parsed.skin
                )

            if parsed.msg_type == proto.ServerMessageType.SPAWNBASELINE:
                baseline_entities[parsed.entity_num] = _EntityInfo(
                    model_num=parsed.model_num,
                    frame=parsed.frame,
                    skin=parsed.skin,
                    origin=parsed.origin,
                    angles=parsed.angles,
                )

            if parsed.msg_type == proto.ServerMessageType.UPDATE:
                baseline = baseline_entities[parsed.entity_num]
                prev_info = entities.get(parsed.entity_num, baseline)
                entities[parsed.entity_num] = prev_info.update(parsed, baseline)
                updated.add(parsed.entity_num)

            if parsed.msg_type == proto.ServerMessageType.TEMP_ENTITY:
                if parsed.temp_entity_type == proto.TempEntityTypes.EXPLOSION:
                    obj_mgr.create_explosion(parsed.origin, time)
                elif parsed.temp_entity_type == proto.TempEntityTypes.TELEPORT:
                    obj_mgr.create_teleport(parsed.origin, time)

            if parsed.msg_type in (
                    proto.ServerMessageType.INTERMISSION,
                    proto.ServerMessageType.FINALE,
                    proto.ServerMessageType.CUTSCENE):
                obj_mgr.set_intermission(True)

        if update_done:
            logger.debug('Handling update. time=%s', time)
            obj_mgr.update(time, prev_entities, entities, prev_updated, updated, fixed_view_angles)
            last_time = time
            prev_updated = updated

    obj_mgr.done(last_time)

    return obj_mgr.world_obj, obj_mgr


import argparse
import logging
import os
import json
import pickle
import random
import time
import numpy as np
import torch
from torch.utils.data import TensorDataset, DataLoader, RandomSampler, SequentialSampler, Dataset
from transformers import AdamW, get_linear_schedule_with_warmup, set_seed
from utils import set_logger, set_seed
from transformers import BertConfig, BertTokenizer
from transformers import AutoModelForSequenceClassification, AutoConfig, AutoTokenizer, AutoModel
from net.bert_base import BertForSequenceClassification
from net.bert_attention import BertForSequenceClassificationAttention
from net.bert_lstm import BertForSequenceClassificationLSTM
from net.bert_lstm_attention import BertForSequenceClassificationLSTMAttenion

from processor import sentiment_processors as processors
from processor import sentiment_convert_examples_to_features, SentimentDataset

from train_and_eval import train, test, predict, _predict

logger = logging.getLogger(__name__)

MODEL_CLASSES = {
    "bert_base": (BertConfig, BertForSequenceClassification, BertTokenizer),
    "bert_attention": (BertConfig, BertForSequenceClassificationAttention, BertTokenizer),
    "bert_lstm": (BertConfig, BertForSequenceClassificationLSTM, BertTokenizer),
    "bert_lstm_attention": (BertConfig, BertForSequenceClassificationLSTMAttenion, BertTokenizer),
}

def main():
    parser = argparse.ArgumentParser()

    # Required parameters
    parser.add_argument("--log_dir", default="log", type=str, required=True, help="设置日志的输出目录")
    parser.add_argument(
        "--dataset",
        choices=["ISEAR", "TEC", "IECE", "SMP2020"],
        default="ISEAR",
        type=str,
        help="应用的数据集，ISEAR, TEC, IECE, SMP2020中4选1",
    )
    parser.add_argument(
        "--model_name",
        default=None,
        type=str,
        required=True,
        help="Model type selected in the list: " + ", ".join(MODEL_CLASSES.keys()),
    )
    parser.add_argument(
        "--pre_train_path",
        default=None,
        type=str,
        required=True,
        help="预训练模型所在的路径，包括 pytorch_model.bin, vocab.txt, bert_config.json",
    )
    parser.add_argument(
        "--output_dir",
        default="output",
        type=str,
        required=True,
        help="The output directory where the model predictions and checkpoints will be written.",
    )
    # Other parameters
    parser.add_argument("--do_train", action="store_true", help="Whether to run training.")
    parser.add_argument("--do_eval", action="store_true", help="Whether to run eval on the dev set.")

    parser.add_argument("--max_seq_length", default=256, type=int, help="输入到bert的最大长度，通常不应该超过512")
    parser.add_argument("--do_lower_case", action="store_true", help="Set this flag if you are using an uncased model.")

    parser.add_argument("--num_train_epochs", default=20, type=int, help="epoch 数目")
    parser.add_argument("--train_batch_size", default=8, type=int, help="训练集的batch_size")
    parser.add_argument("--eval_batch_size", default=512, type=int, help="验证集的batch_size")
    parser.add_argument("--gradient_accumulation_steps", type=int, default=1, help="梯度累计更新的步骤，用来弥补GPU过小的情况")
    parser.add_argument("--learning_rate", default=5e-5, type=float, help="学习率")
    parser.add_argument("--weight_decay", default=0.01, type=float, help="权重衰减")
    parser.add_argument("--adam_epsilon", default=1e-8, type=float, help="Epsilon for Adam optimizer.")
    parser.add_argument("--max_grad_norm", default=1.0, type=float, help="最大的梯度更新")
    parser.add_argument("--seed", type=int, default=233, help="random seed for initialization")
    # parser.add_argument("--warmup_steps", default=0, type=int,
    #                     help="让学习增加到1的步数，在warmup_steps后，再衰减到0")
    parser.add_argument(
        "--warmup_rate", default=0.00, type=float, help="让学习增加到1的步数，在warmup_steps后，再衰减到0，这里设置一个小数，在总训练步数*rate步时开始增加到1"
    )

    args = parser.parse_args()

    args.output_dir = os.path.join(args.output_dir, args.dataset + "_" + args.model_name)
    if not os.path.exists(args.output_dir):
        os.makedirs(args.output_dir)

    assert os.path.exists(os.path.join("data", args.dataset))
    assert os.path.exists(args.pre_train_path)
    assert os.path.exists(args.output_dir)

    # 暂时不写多GPU
    args.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    set_seed(args.seed)

    log_dir = os.path.join(
        args.log_dir,
        args.dataset + "_" + args.model_name + time.strftime("%Y_%m_%d_%H_%M_%S", time.localtime(time.time())) + ".log",
    )
    set_logger(log_dir)

    data_dir = os.path.join("data", args.dataset)

    processor = processors[args.dataset](args, data_dir)
    label_list = processor.get_labels()
    num_labels = len(label_list)

    args.model_name = args.model_name.lower()
    config_class, model_class, tokenizer_class = MODEL_CLASSES[args.model_name]
    if args.do_train:
        logging.info("loading pretrained model... ...")
        config = config_class.from_pretrained(args.pre_train_path, num_labels=num_labels)
        tokenizer = tokenizer_class.from_pretrained(args.pre_train_path, do_lower_case=args.do_lower_case)
        config.save_pretrained(args.output_dir)
        tokenizer.save_vocabulary(args.output_dir)
        model = model_class.from_pretrained(args.pre_train_path, config=config, args=args)
        model.to(args.device)
        logging.info("load pretrained model end... ...")
        logger.info("Training parameters %s", args)

    def convert_to_dataset(examples):
        features = sentiment_convert_examples_to_features(
            examples=examples, tokenizer=tokenizer, max_length=args.max_seq_length, label_list=label_list
        )
        return SentimentDataset(features)

    # Training
    if args.do_train:
        logging.info("loading dataset... ...")
        train_examples = processor.get_train_examples()
        train_dataset = convert_to_dataset(train_examples)
        dev_examples = processor.get_dev_examples()
        dev_dataset = convert_to_dataset(dev_examples)
        logging.info("dataset loaded...")

        train_dataset = np.array(train_dataset)
        dev_dataset = np.array(dev_dataset)

        logging.info("start training... ...")
        train(args, train_dataset, dev_dataset, model)
        logging.info("train end...")

    if args.do_eval:
        logging.info("loading trained model... ...")
        tokenizer = tokenizer_class.from_pretrained(args.output_dir, do_lower_case=args.do_lower_case)
        config = config_class.from_pretrained(args.output_dir, num_labels=num_labels)
        model = model_class.from_pretrained(args.output_dir, config=config, args=args)
        model.to(args.device)
        logging.info("load trained model end... ...")
        logger.info("Evaluation parameters %s", args)

    # Evaluation
    if args.do_eval:
        logging.info("loading dataset... ...")
        test_examples = processor.get_test_examples()
        test_dataset = convert_to_dataset(test_examples)
        logging.info("dataset loaded...")

        test_dataset = np.array(test_dataset)
        test_dataset = np.array(test_dataset)

        logging.info("start evaluating... ...")
        test_probs = test(args, model, test_dataset)
        logging.info("evaluate end...")

if __name__ == "__main__":
    main()


#!/usr/bin/env python2
#
# Strelka - Small Variant Caller
# Copyright (c) 2009-2018 Illumina, Inc.
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.
#
#

"""
Given an EVS model evaluation, provide binned scores and empirical precision
"""

__author__ = "Konrad Scheffler <kscheffler@illumina.com>"

import os
import sys
import random
import numpy
import pandas
import json
from sklearn import linear_model

def parseArgs():
    import argparse

    parser = argparse.ArgumentParser("evs QQ script")
    parser.add_argument("inputs", help="Quality/tag CSV files", nargs="+")
    parser.add_argument("-o", "--output", required=True,
                        help="QQ output file name (csv)")
    parser.add_argument("-c", "--calibration", required=True,
                        help="Calibration output file name (json)")
    parser.add_argument("-N", "--bins", dest="nbins", type=int, default=20,
                        help="Number of bins for which to calculate empirical precision.")
    parser.add_argument("-G", "--geno", dest="geno", type=int, default="-1",
                        help="If specified, restrict analysis to variants with this GENO value.")

    args = parser.parse_args()

    def checkFile(filename, label) :
        if not os.path.isfile(filename) :
            raise Exception("Can't find input %s file: '%s'" % (label,filename))

    if len(args.inputs) == 0 :
        raise Exception("No input file(s) given")

    for inputFile in args.inputs :
        checkFile(inputFile,"features CSV")

    return args

def phred(corrprob):
    return numpy.minimum(-10*numpy.log10(1-corrprob),numpy.full(corrprob.shape,50))

def main():
    args = parseArgs()

    datasets = []
    for i in args.inputs:
        i = os.path.abspath(i)
        print "Reading %s" % i
        df = pandas.read_csv(i)
        datasets.append(df)

    if len(datasets) > 1:
        dataset = pandas.concat(datasets)
    else:
        dataset = datasets[0]

    if args.geno >= 0:
        dataset = dataset[dataset["GENO"] == args.geno]

    result = []
    regression = []

    print "Total: %d" % len(dataset)
    fns = dataset[dataset["tag"] == "FN"].shape[0]
    print "fn: %d" % fns
    print "tp: %d" % len(dataset[dataset["tag"] == "TP"])
    print "fp: %d" % len(dataset[dataset["tag"] == "FP"])
    print "unk: %d" % len(dataset[dataset["tag"] == "UNK"])
    dataset = pandas.DataFrame(dataset[dataset["tag"].isin(["TP", "FP"])])
    data = pandas.DataFrame(dataset.dropna(subset=["qual"]))
    data = data.iloc[numpy.random.permutation(len(data))] # shuffle before sorting: this ensures random order within tied groups
    data = data.sort("qual")

    # note this call fails on numpy 1.9.x (default package for Ubuntu 14.04) TODO: workaround?
    bins = numpy.array_split(data,args.nbins)

    counter = 0
    for bin in bins:
        rec = {"qual": numpy.median(bin["qual"])}
        rec["tp"] = bin[bin["tag"] == "TP"].shape[0]
        rec["fp"] = bin[bin["tag"] == "FP"].shape[0]
        rec["na"] = bin[bin["tag"] == "UNK"].shape[0]

        try:
            rec["precision"] = rec["tp"] / float(rec["tp"] + rec["fp"])
        except:
            rec["precision"] = None
        try:
            rec["fracNA"] = rec["na"] / float(rec["tp"] + rec["fp"] + rec["na"])
        except:
            rec["fracNA"] = None
        result.append(rec)

        counter += 1
        if counter % 10 == 0:
            print "Processed %i / %i values" % (counter, len(bins))

    df = pandas.DataFrame(result)
    df["Q"] = phred(df["qual"])
    df["empiricalQ"] = phred(df["precision"])
    regr = linear_model.LinearRegression()
    regr.fit(phred(df["qual"]).reshape(-1,1), phred(df["precision"]).reshape(-1,1))
    df["calibratedQ"] = regr.predict(phred(df["qual"]).reshape(-1,1))
    print df
    print("Coefficient: {}, Intercept: {}".format(regr.coef_[0,0], regr.intercept_[0]))
    df.to_csv(args.output)
    json.dump({"Coefficient" : regr.coef_[0,0], "Intercept" : regr.intercept_[0]}, open(args.calibration, 'w'))

if __name__ == '__main__':
    main()


#!/usr/bin/env python

import argparse
import random
import numpy as np
import copy
import re
import langdetect
import os.path
from itertools import izip
from collections import defaultdict

from cocoa.core.schema import Schema
from cocoa.core.scenario_db import ScenarioDB
from cocoa.core.util import generate_uuid, write_json, read_json
import cocoa.options

from core.scenario import Scenario
from core.kb import KB

BUYER = Scenario.BUYER
SELLER = Scenario.SELLER

def is_valid_line(line):
    if 'contact' in line.lower():
        return False
    if not re.search(r'\.|\!|\,', line) and len(line.split()) > 15:
        return False
    if re.search(r'\$\s*\d+', line):
        return False
    try:
        if langdetect.detect(line) != 'en':
            return False
    except langdetect.lang_detect_exception.LangDetectException:
        return True
    return True

def process_listing(listing):
    if listing['category'] == 'car' and listing['price'] < 3000:
        return None

    lines = []
    for line in listing['description']:
        if not is_valid_line(line):
            continue
        lines.append(line)

    num_words = sum([len(line.split()) for line in lines])
    if num_words < 20 or num_words > 200:
        return None

    listing['description'] = lines

    return listing

def generate_kbs(schema, listing):
    buyer_item, seller_item = {}, {}
    for attr in schema.attributes:
        if attr.name in ('Role', 'Target', 'Bottomline'):
            continue
        l = listing[attr.name.lower()]
        if attr.name == 'Description':
            # NOTE: Buyer only sees the first half
            N = max(1, len(l) / 2)
            buyer_item[attr.name] = l[:N]
        else:
            buyer_item[attr.name] = l
        seller_item[attr.name] = listing[attr.name.lower()]
    seller_kb = KB(schema.attributes, {'personal': {'Role': 'seller'}, 'item': seller_item})
    buyer_kb = KB(schema.attributes, {'personal': {'Role': 'buyer'}, 'item': buyer_item})
    kbs = [None, None]
    kbs[BUYER] = buyer_kb
    kbs[SELLER] = seller_kb
    return kbs

def discretize(price, price_unit):
    price = int(price / price_unit)
    return price

def generate_price_range(base_price, price_unit, discounts):
    '''
    base: a middle point to generate the range
    intersections: percentage of intersection relative to the range
    '''
    base_price = discretize(base_price, price_unit)
    seller_target = base_price
    for i in discounts:
        buyer_target = i * seller_target
        # Reverse discretization
        seller_target = int(seller_target * price_unit)
        buyer_target = int(buyer_target * price_unit)
        if seller_target == 0 or buyer_target == 0:
            continue
        yield {
                SELLER: {'Bottomline': None, 'Target': seller_target},
                BUYER: {'Bottomline': None, 'Target': buyer_target},
              }

def generate_scenario(schema, base_price, price_unit, discounts, listings):
    for listing in listings:
        listing = process_listing(listing)
        if listing:
            base_price = int(listing['price'])
            if base_price < price_unit:
                continue
            for ranges in generate_price_range(base_price, price_unit, discounts):
                kbs = generate_kbs(schema, listing)
                kbs[BUYER].facts['personal'].update(ranges[BUYER])
                kbs[SELLER].facts['personal'].update(ranges[SELLER])
                yield Scenario(generate_uuid('S'), listing['post_id'], listing['category'], listing['images'], schema.attributes, kbs)

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--random-seed', help='Random seed', type=int, default=1)
    parser.add_argument('--num-scenarios', help='Number of scenarios to generate', type=int, default=-1)
    parser.add_argument('--discounts', nargs='*', type=float, default=[0.9, 0.7, 0.5], help="Discount that the buyer wants to have")
    parser.add_argument('--scraped-data', required=True, help="Path to scraped data")
    parser.add_argument('--categories', nargs='+', required=True, help="Listing categories, JSON file path is craigslist_<category>.json")
    parser.add_argument('--fractions', nargs='+', required=True, help="Fractions of data from different categories")
    parser.add_argument('--skip', default=0, type=int, help="Skip the first X scenarios")
    cocoa.options.add_scenario_arguments(parser)
    args = parser.parse_args()

    random.seed(args.random_seed)
    np.random.seed(args.random_seed)

    schema = Schema(args.schema_path)

    listings = [read_json(os.path.join(args.scraped_data, 'craigslist_{}.json'.format(c))) for c in args.categories]
    fractions = np.array([float(x) for x in args.fractions])
    fractions = fractions / np.sum(fractions)

    # Sample listings
    sampled_listings = []
    N = sum([len(l) for l in listings])
    for listing, fraction in izip(listings, fractions):
        n = int(N * fraction)
        print listing[0]['category'], len(listing), fraction, n
        sampled_listings.append(listing[:n])
    listings = [x for l in sampled_listings for x in l]
    N = len(listings)
    inds = np.random.permutation(N)
    listings = [listings[i] for i in inds]

    base_price = None

    scenario_list = []
    price_unit = 1  # Just use the real price
    scenario_generator = generate_scenario(schema, base_price, price_unit, args.discounts, listings)
    for i, s in enumerate(scenario_generator):
        if i % 100 == 0:
            print i
        if i < args.skip:
            continue
        if len(scenario_list) == args.num_scenarios:
            break
        scenario_list.append(s)
    if len(scenario_list) < args.num_scenarios:
        print 'Not enough listings: {} scenarios generated.'.format(len(scenario_list))
    scenario_db = ScenarioDB(scenario_list)
    write_json(scenario_db.to_dict(), args.scenarios_path)

    for i in range(min(10, len(scenario_db.scenarios_list))):
        print '---------------------------------------------------------------------------------------------'
        print '---------------------------------------------------------------------------------------------'
        scenario = scenario_db.scenarios_list[i]
        print "Scenario id: %s" % scenario.uuid
        for agent in (0, 1):
            kb = scenario.kbs[agent]
            kb.dump()

    num_listings_per_category = defaultdict(int)
    for s in scenario_list:
        cat = s.kbs[0].facts['item']['Category']
        num_listings_per_category[cat] += 1
    for k, v in num_listings_per_category.iteritems():
        print k, v
    print '%d scenarios generated' % len(scenario_list)


import os
import sys
import glob
import time
import copy
import logging
import argparse
import random
import numpy as np
import utils
import tensorflow as tf
import tensorflow.keras as keras
import tensorflow.keras.layers as layers
from nao import NAO

parser = argparse.ArgumentParser()
# Basic model parameters.
parser.add_argument('--data', type=str, default='data')
parser.add_argument('--output_dir', type=str, default='models')
parser.add_argument('--seed', type=int, default=0)
parser.add_argument('--new_arch', type=int, default=300)
parser.add_argument('--encoder_layers', type=int, default=1)
parser.add_argument('--encoder_hidden_size', type=int, default=64)
parser.add_argument('--encoder_emb_size', type=int, default=32)
parser.add_argument('--mlp_layers', type=int, default=0)
parser.add_argument('--mlp_hidden_size', type=int, default=200)
parser.add_argument('--decoder_layers', type=int, default=1)
parser.add_argument('--decoder_hidden_size', type=int, default=64)
parser.add_argument('--source_length', type=int, default=40)
parser.add_argument('--encoder_length', type=int, default=20)
parser.add_argument('--decoder_length', type=int, default=40)
parser.add_argument('--encoder_dropout', type=float, default=0)
parser.add_argument('--mlp_dropout', type=float, default=0)
parser.add_argument('--decoder_dropout', type=float, default=0)
parser.add_argument('--l2_reg', type=float, default=0.0)
parser.add_argument('--encoder_vocab_size', type=int, default=12)
parser.add_argument('--decoder_vocab_size', type=int, default=12)
parser.add_argument('--max_step_size', type=int, default=100)
parser.add_argument('--trade_off', type=float, default=0.8)
parser.add_argument('--epochs', type=int, default=1000)
parser.add_argument('--batch_size', type=int, default=100)
parser.add_argument('--lr', type=float, default=0.001)
parser.add_argument('--optimizer', type=str, default='adam')
parser.add_argument('--grad_bound', type=float, default=5.0)
parser.add_argument('--expand', type=int, default=None)
parser.add_argument('--iteration', type=int, default=0)
parser.add_argument('--generate_topk', type=int, default=100)
parser.add_argument('--remain_topk', type=int, default=100)
args = parser.parse_args()

log_format = '%(asctime)s %(message)s'
logging.basicConfig(stream=sys.stdout, level=logging.INFO,
    format=log_format, datefmt='%m/%d %I:%M:%S %p')

def nao_train(train_queue, model, optimizer):
    objs = utils.AvgMeter()
    mse = utils.AvgMeter()
    nll = utils.AvgMeter()
    model.train()
    for step, sample in enumerate(train_queue):
        encoder_input = sample['encoder_input']
        encoder_target = sample['encoder_target']
        decoder_input = sample['decoder_input']
        decoder_target = sample['decoder_target']
        with tf.GradientTape() as tape:
            predict_value, log_prob, arch = model(encoder_input, decoder_input, training=True)
            loss_1 = keras.losses.MSE(encoder_target, predict_value)
            loss_2 = keras.losses.categorical_crossentropy(decoder_target, log_prob)
            loss = args.trade_off * loss_1 + (1 - args.trade_off) * loss_2
        gradients = tape.gradient(loss, model.trainable_variables)
        if args.grad_bound != 0.0:
            gradients, _ = tf.clip_by_global_norm(gradients, args.grad_bound)
        optimizer.apply_gradients(zip(gradients, model.trainable_variables))

        n = encoder_input.shape[0]
        objs.update(loss, n)
        mse.update(loss_1, n)
        nll.update(loss_2, n)

    return objs.avg, mse.avg, nll.avg

def nao_valid(queue, model):
    pa = utils.AvgMeter()
    hs = utils.AvgMeter()
    mse = utils.AvgMeter()

    for step, sample in enumerate(queue):
        encoder_input = sample['encoder_input']
        encoder_target = sample['encoder_target']
        decoder_target = sample['decoder_target']

        predict_value, logits, arch = model(encoder_input)
        n = encoder_input.shape[0]
        pairwise_acc = utils.pairwise_accuracy(encoder_target.data.squeeze().tolist(),
                                               predict_value.data.squeeze().tolist())
        hamming_dis = utils.hamming_distance(decoder_target.data.squeeze().tolist(), arch.data.squeeze().tolist())
        mse.update(keras.losses.MSE(encoder_target.data.squeeze(), predict_value.data.squeeze()), n)
        pa.update(pairwise_acc, n)
        hs.update(hamming_dis, n)

    return mse.avg, pa.avg, hs.avg

def nao_infer(queue, model, step, direction=1): #direction='+'
    new_arch_list = []
    model.eval()
    for i, sample in enumerate(queue):
        encoder_input = sample['encoder_input']
        new_arch = model.generate_new_arch(encoder_input, step, direction=direction)
        new_arch_list.extend(new_arch.data.squeeze().tolist())
    return new_arch_list

def main():
    random.seed(args.seed)
    np.random.seed(args.seed)
    tf.random.set_seed(args.seed)

    logging.info("Args = %s", args)

    nao = NAO(
        args.encoder_layers,
        args.encoder_vocab_size,
        args.encoder_hidden_size,
        args.encoder_dropout,
        args.encoder_length,
        args.source_length,
        args.encoder_emb_size,
        args.mlp_layers,
        args.mlp_hidden_size,
        args.mlp_dropout,
        args.decoder_layers,
        args.decoder_vocab_size,
        args.decoder_hidden_size,
        args.decoder_dropout,
        args.decoder_length,
    )
    logging.info("param size = %fMB", utils.count_parameters_in_MB(nao))
    nao = nao.cuda()

    with open(os.path.join(args.output_dir, 'arch_pool.{}'.format(args.iteration))) as f:
        arch_pool = f.read().splitlines()
        arch_pool = list(map(utils.build_dag, arch_pool))
    with open(os.path.join(args.output_dir, 'arch_pool.{}.perf'.format(args.iteration))) as f:
        arch_pool_valid_acc = f.read().splitlines()
        arch_pool_valid_acc = list(map(float, arch_pool_valid_acc))

    logging.info('Training Encoder-Predictor-Decoder')
    train_encoder_input = list(
        map(lambda x: utils.parse_arch_to_seq(x[0], 2) + utils.parse_arch_to_seq(x[1], 2), arch_pool))
    min_val = min(arch_pool_valid_acc)
    max_val = max(arch_pool_valid_acc)
    train_encoder_target = [(i - min_val) / (max_val - min_val) for i in arch_pool_valid_acc]

    if args.expand is not None:
        buffer1, buffer2 = [], []
        for _ in range(args.expand - 1):
            for src, tgt in zip(train_encoder_input, train_encoder_target):
                a = np.random.randint(0, 5)
                b = np.random.randint(0, 5)
                src = src[:4 * a] + src[4 * a + 2:4 * a + 4] + \
                      src[4 * a:4 * a + 2] + src[4 * (a + 1):20 + 4 * b] + \
                      src[20 + 4 * b + 2:20 + 4 * b + 4] + src[20 + 4 * b:20 + 4 * b + 2] + \
                      src[20 + 4 * (b + 1):]
                buffer1.append(src)
                buffer2.append(tgt)
        train_encoder_input += buffer1
        train_encoder_target += buffer2

    nao_train_dataset = utils.NAODataset(train_encoder_input, train_encoder_target, True, swap=True)
    nao_valid_dataset = utils.NAODataset(train_encoder_input, train_encoder_target, False)
    nao_train_queue = torch.utils.data.DataLoader(
        nao_train_dataset, batch_size=args.batch_size, shuffle=True, pin_memory=True)
    nao_valid_queue = torch.utils.data.DataLoader(
        nao_valid_dataset, batch_size=len(nao_valid_dataset), shuffle=False, pin_memory=True)
    nao_optimizer = torch.optim.Adam(nao.parameters(), lr=args.lr, weight_decay=args.l2_reg)
    for nao_epoch in range(1, args.epochs + 1):
        nao_loss, nao_mse, nao_ce = nao_train(nao_train_queue, nao, nao_optimizer)
        if nao_epoch % 10 == 0 or nao_epoch == 1:
            logging.info("epoch %04d train loss %.6f mse %.6f ce %.6f", nao_epoch, nao_loss, nao_mse, nao_ce)
        if nao_epoch % 100 == 0 or nao_epoch == 1:
            mse, pa, hs = nao_valid(nao_train_queue, nao)
            logging.info("Evaluation on train data")
            logging.info('epoch %04d mse %.6f pairwise accuracy %.6f hamming distance %.6f', nao_epoch, mse, pa, hs)
            mse, pa, hs = nao_valid(nao_valid_queue, nao)
            logging.info("Evaluation on valid data")
            logging.info('epoch %04d mse %.6f pairwise accuracy %.6f hamming distance %.6f', nao_epoch, mse, pa, hs)

    new_archs = []
    predict_step_size = 0
    top_archs = list(map(lambda x: utils.parse_arch_to_seq(x[0], 2) + utils.parse_arch_to_seq(x[1], 2),
                         arch_pool[:args.generate_topk]))
    nao_infer_dataset = utils.NAODataset(top_archs, None, False)
    nao_infer_queue = torch.utils.data.DataLoader(nao_infer_dataset, batch_size=len(nao_infer_dataset), shuffle=False,
                                                  pin_memory=True)

    while len(new_archs) < args.new_arch:
        predict_step_size += 1
        logging.info('Generate new architectures with step size %d', predict_step_size)
        new_arch = nao_infer(nao_infer_queue, nao, predict_step_size, direction='+')
        for arch in new_arch:
            if arch not in train_encoder_input and arch not in new_archs:
                new_archs.append(arch)
            if len(new_archs) >= args.new_arch:
                break
        logging.info('%d new archs generated now', len(new_archs))
        if predict_step_size > args.max_step_size:
            break

    logging.info("Generate %d new archs", len(new_archs))
    new_arch_pool = list(map(lambda x: utils.parse_seq_to_arch(x, 2), new_archs))
    new_arch_pool = new_arch_pool + arch_pool[:args.remain_topk]
    with open(os.path.join(args.output_dir, 'new_arch_pool.{}'.format(args.iteration)), 'w') as f:
        for arch in new_arch_pool:
            arch = ' '.join(map(str, arch[0] + arch[1]))
            f.write('{}\n'.format(arch))
    logging.info('Finish training!')

if __name__ == '__main__':
    main()

def train_nao(model,
              env,
              nao_dataset,
              search_space,
              epochs,
              optimizer,
              alpha):

    acc_buffer = []
    pred_acc_buff = []
    ema_reward_buffer = []
    loss_buffer = []

    for epoch in range(epochs):
        print('\n')
        print('-' * 40, "epoch %d:" % (epoch + 1), '-' * 40)

        encoder_input_tokens, decoder_input_tokens = nao_dataset
        acc = env.get_acc(encoder_input_tokens)
        labels = tf.one_hot(encoder_input_tokens, axis=-1)
        with tf.GradientTape() as tape:
            pred_acc, all_logits, all_pred_tokens = model(encoder_input_tokens, decoder_input_tokens)
            reconstruct_loss = tf.nn.softmax_cross_entropy_with_logits(labels, all_logits, axis=-1, name=None)
            prediction_loss = (acc - pred_acc)**2
            total_loss = alpha * prediction_loss + (1 - alpha) * reconstruct_loss

        gradients = tape.gradient(total_loss, model.trainable_variables)
        # clip and apply gradient
        # if clip_norm != 0.0:
        #     gradients, _ = tf.clip_by_global_norm(gradients, clip_norm)
        # print(gradients)
        optimizer.apply_gradients(zip(gradients, model.trainable_variables))

        acc_buffer.append(acc)
        pred_acc_buff.append(pred_acc)
        loss_buffer.append(total_loss.numpy())

    return acc_buffer, pred_acc_buff, loss_buffer


#!/usr/bin/env python
from __future__ import division

import math
import numpy as np
import nltk

def wrapGenTok(tok):
    #wrap tok inside `GenToken[]`
    return 'GenToken[' + tok + ']'

def read_corpus(data_file):
    """
    extract input tokens using NLTK, whereas leave output tokens unformatted
    @param data_file (str): file containing src_sent and tgt_sent sep by '\t', each line is delineated by '\n'
    @return data ((list[list[str]], list[str])): tuples of list of src and tgt tokens
    """
    src_data, tgt_data = [], []
    for sent in open(data_file, 'r'):
        sents = sent.split('\n')[0].split('\t') #[src_sent, tgt_sent]
        src_sent = sents[0]
        tgt_sent = sents[1]

        src_sent_toks = nltk.word_tokenize(src_sent)
        src_data.append(src_sent_toks)

        tgt_data.append(tgt_sent)

    return src_data, tgt_data

def pad_sents(sents, pad_id=0):
    """
    pad the list of sents according to max sent len
    @param sents (list[list[int]]): list of word ids of sentences
    @param pad_id (int): pad idx
    @return sents_padded (list[list[int]]): padded sentences
    """
    sents_padded = []
    max_len = 0

    for sent in sents:
        if len(sent) > max_len: max_len = len(sent)
    for sent in sents:
        sent_padded = sent
        sent_padded.extend([pad_id for i in range(max_len - len(sent))])
        sents_padded.append(sent_padded)

    return sents_padded

def save_sents(sents, file_path):
    """
    save sentences in a file line by line
    @param sents (list[str]): list of sentences
    @param file_path (str): location to save senetnces
    """
    with open(file_path, 'w') as file_obj:
        for sent in sents:
            file_obj.write(sent + '\n')

def batch_iter(src_sents, tgt_sents, lang, batch_size, shuffle=False):
    """
    Yield batches of source sentences and target nodes and actions lengths
    @param src_sents (list(list[str])): list of source sentences (list of tokens)
    @param tgt_sents (list[str]): list of target sentences
    @param lang: target language
    @param batch_size (int): batch size
    @param shuffle (boolean): whether to randomly shuffle the dataset
    """
    if lang == 'lambda':
        from lang.Lambda.lambda_transition_system import LambdaCalculusTransitionSystem
        from lang.Lambda.transition_system import ApplyRuleAction, ReduceAction, GenTokenAction
        from lang.Lambda.asdl import ASDLGrammar
        from lang.Lambda.parse import parse_lambda_expr, logical_form_to_ast

        asdl_desc = open('lang/Lambda/lambda_asdl.txt').read()
        grammar = ASDLGrammar.from_text(asdl_desc)
        parser = LambdaCalculusTransitionSystem(grammar)

        tgt_asts = [logical_form_to_ast(grammar, parse_lambda_expr(sent)) for sent in tgt_sents]
        tgt_rules = [parser.get_actions(ast) for ast in tgt_asts]

        tgt_nodes, tgt_tokens, tgt_actions = [], [], []
        for actions in tgt_rules:
            tgt_node, tgt_token, tgt_action = [], [], []
            nodes = ['<start>']
            for action in actions:
                assert len(nodes) > 0
                #only include token if last action was GenToken
                if len(tgt_action) > 0 and tgt_action[-1] != 'Reduce' and type(tgt_action[-1]) == str:
                    tgt_token.append(tgt_action[-1])
                else:
                    tgt_token.append('<pad>')
                if isinstance(action, ApplyRuleAction):
                    if grammar.mul_cardinality(nodes[-1]):
                        tgt_node.append(nodes[-1])
                    else:
                        tgt_node.append(nodes.pop())

                    rule = action.production #ASDLProduction(ASDLType(type), ASDLConstructor(constructor))
                    fields = rule.constructor.fields #fields[Field(name, type, cardinality)]
                    curr_nodes = []
                    for field in fields: #Field(name, ASDLType(name), cardinality)
                        node_name = field.type.name
                        field_cardinality = field.cardinality
                        if field_cardinality == 'multiple':
                            node_name += '*'
                        curr_nodes.append(node_name)
                    nodes.extend(curr_nodes[::-1])

                    tgt_action.append(rule)
                elif isinstance(action, GenTokenAction):
                    tgt_node.append(nodes.pop())
                    tgt_action.append(action.token)
                else: #ReduceAction
                    tgt_node.append(nodes.pop())
                    tgt_action.append('Reduce')
            tgt_nodes.append(tgt_node)
            tgt_tokens.append(tgt_token)
            tgt_actions.append(tgt_action)

        data = zip(src_sents, tgt_nodes, tgt_tokens, tgt_actions)

        batch_num = int(math.ceil(len(data) / batch_size))
        index_array = list(range(len(data)))

        if shuffle:
            np.random.shuffle(index_array)

        for i in range(batch_num):
            indices = index_array[i * batch_size: (i + 1) * batch_size]
            examples = [data[idx] for idx in indices]

            examples = sorted(examples, key=lambda e: len(e[0]), reverse=True)
            batch_sents = [e[0] for e in examples]
            batch_nodes = [e[1] for e in examples]
            batch_tokens = [e[2] for e in examples]
            batch_actions = [e[3] for e in examples]

            yield batch_sents, batch_nodes, batch_tokens, batch_actions
    else:
        print('language:  %s currently not supported' % (lang))

def comp_exact_match(refs, hyps):
    """
    @param refs (list[list[str]]): list of reference sentences
    @param hyps (list[list[str]]): list of predicted sentences (hypotheses)
    @return em (float): average exact match value
    """
    assert len(refs) == len(hyps)
    match = 0
    for ref, hyp in zip(refs, hyps):
        if ref == hyp: match += 1
    return match / len(refs)


import os
from os.path import isfile, join
import re
import numpy as np

def get_phoneme_label_vector(phn_file, length, fs, befaft, mode='phonemes', return_label_lists=False):
    '''
    Creates a time-series vector of phoneme labels based on the output of the Penn-Phonetic Alignment
    procedure.

    Returns label vector as numpy array of shape (time, ) containing categorical labels.
    Assumes that the .phn files were made on wav files that did not have a befaft zero period.

    This function can be used instead of the ``get_label_vecs_from_files`` method
    inside the ``Aligner`` class if you only want a single file's label vector.

    Parameters
    ----------
    phn_file : string, path to .phn file
    length : int, length of stimulus (in samples) including befaft periods
    fs : int, sampling rate
    befaft : array or list of length 2, containing befaft time periods in seconds
    mode : string, either 'phonemes', or 'manner'. Any empty or 'sp' periods are labeled -1.
        If 'manner': labels correspond to ['plosive','fricative','nasal','sonorant']
    return_label_lists : bool, if True, returns a 2-tuple containing the labels as well as the list
        of possible labels, so that the index of a label in the list is the integer label assigned to it.

    Returns
    -------
    labels : np.array, shape (time,)
        Integer labels over time. Full length is befaft[0]*fs+length+befaft[1]*fs

    item_list : list
        List of phonemes (if mode=='phonemes') or manners of articulation (if mode=='manner')
        which were used. The index of a phoneme indicates the label it is assigned.
    '''

    pattern = r'[0-9]'
    phoneme_list = ['EH','K','S','L','AH','M','EY','SH','N','P','OY','T','OW','Z','W','D','B','V','IH','AA','R','AY','ER','AE','AO','NG','G','TH','IY','F','DH','HH','UH','CH','UW','AW','JH','Y','ZH']
    manner_list = ['plosive','fricative','nasal','sonorant']
    manner_dict = {'EH':'sonorant',
                    'K':'plosive',
                    'S':'fricative',
                    'L':'sonorant',
                    'AH':'sonorant',
                    'M':'nasal',
                    'EY':'sonorant',
                    'SH':'fricative',
                    'N':'nasal',
                    'P':'plosive',
                    'OY':'sonorant',
                    'T':'plosive',
                    'OW':'sonorant',
                    'Z':'fricative',
                    'W':'sonorant',
                    'D':'plosive',
                    'B':'plosive',
                    'V':'fricative',
                    'IH':'sonorant',
                    'AA':'sonorant',
                    'R':'sonorant',
                    'AY':'sonorant',
                    'ER':'sonorant',
                    'AE':'sonorant',
                    'AO':'sonorant',
                    'NG':'nasal',
                    'G':'plosive',
                    'TH':'fricative',
                    'IY':'sonorant',
                    'F':'fricative',
                    'DH':'fricative',
                    'HH':'fricative',
                    'UH':'sonorant',
                    'CH':'velar',
                    'UW':'sonorant',
                    'AW':'sonorant',
                    'JH':'velar',
                    'Y':'sonorant',
                    'ZH':'sonorant',
                    'sp':'none'}

    f = open(phn_file, 'r')
    Lines = [line.strip() for line in f.readlines()]
    f.close()

    labels = -1 * np.ones((length,))

    for line in Lines:
        phn = re.sub(pattern, '', line.split(' ')[-1])
        start_time, end_time = line.split(' ')[:2]
        start_time, end_time = float(start_time), float(end_time)
        start_time += befaft[0]
        end_time += befaft[0]
        if mode == 'phonemes':
            labels[round(start_time*fs):round(end_time*fs)] = phoneme_list.index(phn) if phn in phoneme_list else -1.0
        elif mode == 'manner':
            labels[round(start_time*fs):round(end_time*fs)] = manner_list.index(manner_dict[phn]) if manner_dict[phn] in manner_list else -1.0
        else:
            raise Exception(f"mode parameter must be either 'phonemes' or 'manner', but found {mode}")

    if return_label_lists:
        if mode == 'phonemes':
            return labels, phoneme_list
        elif mode == 'manner':
            return labels, manner_list
        else:
            raise Exception(f"mode parameter must be either 'phonemes' or 'manner', but found {mode}")
    return labels

def get_word_label_vector(wrd_file, length, fs, befaft, wrd_dict=None, wrd_files_dir=None, return_wrd_dict=False):
    '''
    Returns label vector as numpy array of shape (time, ) containing categorical labels.
    Assumes that the .wrd files were made on wav files that did not have a befaft zero period.

    This function can be used instead of the ``get_label_vecs_from_files`` method
    inside the ``Aligner`` class if you only want a single file's label vector.

    Parameters
    ----------
    wrd_file : string, path to .wrd file
    length : int
        length of stimulus (in samples) including befaft periods
    fs : int, sampling rate
    befaft : array or list of length 2
        containing befaft time periods in seconds
    wrd_dict : dict
        keys are words (capitalized) and values are integers which become the labels for each word
    wrd_files_dir : string
        Path to directory containing all the .wrd files to be included in the dictionary. This is
        ignored if wrd_dict is supplied, but otherwise, a new wrd_dict is created using the .wrd
        files in this directory. By default, this will make 'sp' (space) have a label of -1.
    return_wrd_dict : bool
        Whether or not to return the wrd_dict along with the labels as a tuple

    Returns
    -------
    labels : np.array, shape (time,)
        Integer labels over time. Full length is befaft[0]*fs+length+befaft[1]*fs
    wrd_dict : dict
        Word dictionary (word to integer) used to create the labels.
        Only returned if return_wrd_dict==True
    '''

    pattern = r'[0-9]'

    if wrd_dict is None and wrd_files_dir is None:
        raise Exception('Must provide either wrd_dict or wrd_files_dir')
    elif wrd_dict is None:
        wrd_dict = create_wrd_dict(wrd_files_dir)

    f = open(wrd_file, 'r')
    Lines = [line.strip() for line in f.readlines()]
    f.close()

    labels = -1 * np.ones((length,))

    for line in Lines:
        wrd = re.sub(pattern, '', line.split(' ')[-1])
        start_time, end_time = line.split(' ')[:2]
        start_time, end_time = float(start_time), float(end_time)
        start_time += befaft[0]
        end_time += befaft[0]
        try:
            labels[round(start_time*fs):round(end_time*fs)] = wrd_dict[wrd]

        except KeyError:
            raise KeyError(f"{wrd} is not a valid key in wrd_dict")

    if return_wrd_dict:
        return labels, wrd_dict

    return labels

def create_wrd_dict(wrd_files_dir, list_to_skip=[]):
    '''
    Create a new word to label dictionary, which can be passed to get_word_label_vector.

    Parameters
    ----------
    wrd_files_dir : string
        Path to directory containing all the .wrd files to be used
    list_to_skip : list of strings, default=[]
        Words which will not be added to the dictionary.

    Returns
    -------
    wrd_dict : dict
        Dictionary of word:int (key:value) pairs for all the words in the corpus
        of files in the directory.
    '''

    pattern = r'[0-9]'

    wrd_files = sorted([f for f in os.listdir(wrd_files_dir) if isfile(join(wrd_files_dir, f)) and f.endswith('.wrd')])

    wrd_dict = {}

    for wrd_file in wrd_files:
        f = open(join(wrd_files_dir, wrd_file), 'r')
        Lines = [line.strip() for line in f.readlines()]
        f.close()

        for line in Lines:
            wrd = re.sub(pattern, '', line.split(' ')[-1])
            if wrd == 'sp':
                wrd_dict[wrd] = -1
            elif wrd not in wrd_dict and wrd not in list_to_skip:
                wrd_dict[wrd] = len(wrd_dict)

    return wrd_dict


import pandas as pd
import numpy as np
from random import randint
import matplotlib.pyplot as plt
import matplotlib.dates as mdates
#import wrapped_flappy_bird as game
import GameEngine as game
import csv

FINAL_EPSILON = 0.0001 # final value of epsilon
INITIAL_EPSILON = 0.1 # starting value of epsilon

if random.random() <= epsilon:
    print("----------Random Action----------")
    action_index = random.randrange(ACTIONS)
    a_t[action_index] = 1

THREADS = 1

game_state = []

for i in range(0,THREADS):
    print("init THREAD")
    game_state.append(game.PlayGame())

games = 10000

actions = [[1,0],[0,1]]

x = []
y = []

for game in range(games):
    terminal = False
    done = False

    #for step in range(gameLength):
    while terminal == False:
        a_t = actions[np.random.choice(np.arange(0, 2), p=[0.5, 0.5])]
        x_t, r_t, terminal = game_state[0].nextStep(a_t)

    if terminal == True:
        profit =  game_state[0].fullBalance -  game_state[0].initialBalance
        reward =  game_state[0].reward
        print("Game:", game, "Profit:", profit, "Reward:", reward)

        y.append(profit)
        x.append( game_state[0].startDate)
        df = pd.DataFrame({"hours": x, "profits": y})
        df.to_csv("aiTestLog.csv")

        game_state[0].startGame()

'''
df = pd.DataFrame({ "hours" : x, "profits" : y })
df["hours"] = pd.to_datetime((df["hours"]))
df = df.set_index("hours")

fig, ax = plt.subplots()
ax.plot(df.index, df["profits"], ".", markersize=1 )

ax.xaxis.set_major_locator(mdates.MonthLocator(interval=1))
ax.xaxis.set_major_formatter(mdates.DateFormatter("%Y-%m"))

fig.autofmt_xdate()
plt.show()
'''


import tensorflow as tf
import numpy as np
import util, model, style_transfer
import os

import argparse

def parse_args():
    desc = "Tensorflow implementation of 'Image Style Transfer Using Convolutional Neural Networks"
    parser = argparse.ArgumentParser(description=desc)

    parser.add_argument('--model_path', type=str, default='.', help='The directory where the pre-trained model was saved')
    parser.add_argument('--content', type=str, default='images/tubingen.jpg', help='File path of content image (notation in the paper : p)', required = True)
    parser.add_argument('--style', type=str, default='images/starry-night.jpg', help='File path of style image (notation in the paper : a)', required = True)
    parser.add_argument('--output', type=str, default='result.jpg', help='File path of output image', required = True)

    parser.add_argument('--loss_ratio', type=float, default=1e-3, help='Weight of content-loss relative to style-loss')

    parser.add_argument('--content_layers', nargs='+', type=str, default=['conv4_2'], help='VGG19 layers used for content loss')
    parser.add_argument('--style_layers', nargs='+', type=str, default=['relu1_1', 'relu2_1', 'relu3_1', 'relu4_1', 'relu5_1'],
                        help='VGG19 layers used for style loss')

    parser.add_argument('--content_layer_weights', nargs='+', type=float, default=[1.0], help='Content loss for each content is multiplied by corresponding weight')
    parser.add_argument('--style_layer_weights', nargs='+', type=float, default=[.2,.2,.2,.2,.2],
                        help='Style loss for each content is multiplied by corresponding weight')

    parser.add_argument('--initial_type', type=str, default='content', choices=['random','content','style'], help='The initial image for optimization (notation in the paper : x)')
    parser.add_argument('--max_size', type=int, default=512, help='The maximum width or height of input images')
    parser.add_argument('--content_loss_norm_type', type=int, default=3, choices=[1,2,3], help='Different types of normalization for content loss')
    parser.add_argument('--num_iter', type=int, default=2, help='The number of iterations to run')

    return check_args(parser.parse_args())

def check_args(args):
    try:
        assert len(args.content_layers) == len(args.content_layer_weights)
    except:
        print ('content layer info and weight info must be matched')
        return None
    try:
        assert len(args.style_layers) == len(args.style_layer_weights)
    except:
        print('style layer info and weight info must be matched')
        return None

    try:
        assert args.max_size > 100
    except:
        print('Too small size')
        return None

    model_file_path = args.model_path + '/' + model.MODEL_FILE
    try:
        assert os.path.exists(model_file_path)
    except:
        print ('There is no %s'%model_file_path)
        return None

    try:
        size_in_KB = os.path.getsize(model_file_path)
        assert abs(size_in_KB - 534904783) < 10
    except:
        print('check file size of \'imagenet-vgg-verydeep-19.mat\'')
        print('there are some files with the same name')
        print('pre_trained_model used here can be downloaded from bellow')
        print('http://www.vlfeat.org/matconvnet/models/imagenet-vgg-verydeep-19.mat')
        return None

    try:
        assert os.path.exists(args.content)
    except:
        print('There is no %s'%args.content)
        return None

    try:
        assert os.path.exists(args.style)
    except:
        print('There is no %s' % args.style)
        return None

    return args

def add_one_dim(image):
    shape = (1,) + image.shape
    return np.reshape(image, shape)

def main():

    # parse arguments
    args = parse_args()
    if args is None:
        exit()

    # initiate VGG19 model
    model_file_path = args.model_path + '/' + model.MODEL_FILE
    vgg_net = model.VGG19(model_file_path)

    # load content image and style image
    content_image = util.load_image(args.content, max_size=args.max_size)
    style_image = util.load_image(args.style, shape=(content_image.shape[1],content_image.shape[0]))

    # initial guess for output
    if args.initial_type == 'content':
        init_image = content_image
    elif args.initial_type == 'style':
        init_image = style_image
    elif args.initial_type == 'random':
        init_image = np.random.normal(size=content_image.shape, scale=np.std(content_image))

    # check input images for style-transfer
    # utils.plot_images(content_image,style_image, init_image)

    # create a map for content layers info
    CONTENT_LAYERS = {}
    for layer, weight in zip(args.content_layers,args.content_layer_weights):
        CONTENT_LAYERS[layer] = weight

    # create a map for style layers info
    STYLE_LAYERS = {}
    for layer, weight in zip(args.style_layers, args.style_layer_weights):
        STYLE_LAYERS[layer] = weight

    # open session
    sess = tf.Session(config=tf.ConfigProto(allow_soft_placement=True))

    # build the graph
    st = style_transfer.StyleTransfer(session = sess,
                                      content_layer_ids = CONTENT_LAYERS,
                                      style_layer_ids = STYLE_LAYERS,
                                      init_image = add_one_dim(init_image),
                                      content_image = add_one_dim(content_image),
                                      style_image = add_one_dim(style_image),
                                      net = vgg_net,
                                      num_iter = args.num_iter,
                                      loss_ratio = args.loss_ratio,
                                      content_loss_norm_type = args.content_loss_norm_type,
                                      )
    # launch the graph in a session
    result_image = st.update()

    # remove batch dimension
    shape = result_image.shape
    result_image = np.reshape(result_image,shape[1:])

    # save result
    util.save_image(result_image,args.output)

if __name__ == '__main__':
    main()