import argparse
import logging
import shutil
import subprocess
import sys
import tempfile
from collections import defaultdict
from dataclasses import dataclass
from pathlib import Path
from typing import List

from func_adl_types_atlas.build_xaod_edm import build_xaod_edm  # type: ignore
from rich.console import Console
from rich.table import Table

from func_adl_servicex_type_generator import generate_package

valid_tests = ["jets_uncalib", "jets_calib", "met", "error_bad_argument"]

def run_command(cmd: str | List[str]):
    """Run a command and dump the resulting output to logging,
    depending on the return code.

    Args:
        cmd (str): Command to run
    """
    if not isinstance(cmd, list):
        cmd = [cmd]

    command_line = ";".join(cmd) if len(cmd) > 1 else f"{cmd[0]}"
    logging.debug(f"Running command: {command_line}")

    try:
        result = subprocess.run(
            command_line,
            capture_output=True,
            text=True,
        )

        level = logging.INFO
        if result.returncode != 0:
            level = logging.ERROR

        logging.log(level, f"Result code from running {cmd}: {result.returncode}")
        for line in result.stdout.split("\n"):
            logging.log(level, f"stdout: {line}")
        for line in result.stderr.split("\n"):
            logging.log(level, f"stderr: {line}")

        if result.returncode != 0:
            raise RuntimeError(
                f"Failed to execute (return code {result.returncode}): {command_line}"
            )
    except Exception:
        logging.error("Exception was thrown running the following commands:")
        for ln in cmd:
            logging.error(f"  {ln}")
        raise

def create_type_json(
    release: str, clean: bool, location: Path, cmd_location: Path
) -> Path:
    """Create the type json file for a release.

    Args:
        release (_type_): Release name - an image tag
        clean (_type_): If true, then refresh even if yaml file is present.
    """
    logging.info(f"Building JSON file for {release}")

    # Can we skip?
    yaml_path = location / f"{release}.yaml"
    if yaml_path.exists() and (not clean):
        logging.debug(f"YAML type file for {release} already exists. Not rebuilding")
        return yaml_path

    # Do the build.
    yaml_path.parent.mkdir(parents=True, exist_ok=True)
    logging.debug(f"Running container to build json type file for {release}")
    build_xaod_edm(release, yaml_path)
    logging.debug(f"Finished building json type file for {release}")

    return yaml_path

def create_python_package(
    release: str, package_version: str, clean: bool, json_location: Path, location: Path
) -> Path:
    """Given the type json file already exists, create the python
    package for the full type information.

    Args:
        release (str): The name of the release
        clean (bool): If true, remake the package from scratch.
                    Otherwise, only if it isn't there already.
        json_location (Path): Location of the json yaml file
        location (Path): Directory to house the created package

    Returns:
        Path: Location of the package
    """
    logging.info(f"Creating python package for release {release}")

    # See if we can bail out quickly
    package_location = location / release
    if package_location.exists() and (not clean):
        logging.debug(
            f"Python package for release {release} already exists. Not rebuilding."
        )
        return package_location

    # Re-create it.
    generate_package(json_location, package_version, package_location)

    # Next we need to find the package - the actual name of the folder might vary a
    # little bit due to the release series.

    return package_location

def do_build_for_release(release, args):
    yaml_location = create_type_json(
        release, args.clean, args.type_json, args.command_location
    )
    return create_python_package(
        release, args.package_version, args.clean, yaml_location, args.type_package
    )

def do_build(args):
    """Iterator that builds a package for each release.

    Args:
        args (): Command line arguments

    Yields:
        Path: THe path where the package is located
    """
    for r in args.release:
        do_build_for_release(r, args)
    return 0

def do_yaml(args):
    for r in args.release:
        create_type_json(r, args.clean, args.type_json, args.command_location)

def do_test(args):
    """After making sure that the packages are built, run the requested tests

    Args:
        args (): Command line arguments
    """
    for r in args.release:
        package_path = do_build_for_release(r, args)
        logging.debug(f"Running tests for release {r} in {package_path}")

        # Build the commands to create the env and setup/run the test.
        commands = []

        test_packages_path = Path(__file__).parent / "test_packages.py"
        assert (
            test_packages_path.exists()
        ), f"Configuration error: cannot find {test_packages_path}"

        logging.debug("About to create the temp directory")
        with tempfile.TemporaryDirectory() as release_dir_tmp:
            release_dir = release_dir_tmp if args.test_dir is None else args.test_dir

            commands.append(f"cd {release_dir}")
            commands.append("python -m venv .venv")
            commands.append(". .venv/Scripts/Activate.ps1")
            commands.append("python -m pip install --upgrade pip")

            # Install the package. We want to run locally, so need to install
            # a special flavor of func_adl_servicex.
            commands.append("python -m pip install func_adl_servicex[local]")
            commands.append(f"python -m pip install -e {package_path.absolute()}")

            # Copy over the script to make it easy to "run".
            shutil.copy(test_packages_path.absolute(), release_dir)

            # Commands to run the scripts
            all_tests = args.test if len(args.test) > 0 else valid_tests
            test_args = " ".join([f"--test {t}" for t in all_tests])
            verbose_arg = (" -" + "v" * int(args.verbose)) if args.verbose > 0 else ""

            commands.append(f"python test_packages.py {test_args}{verbose_arg}")

            # Finally, run the commands.
            logging.debug(f"Running tests for release {r} - commands are:")
            for cmd in commands:
                logging.debug(f"  {cmd}")
            run_command(commands)
    return 0

@dataclass
class ReleaseInfo:
    """Information about a release"""

    has_type_json: bool = False
    has_package: bool = False

def do_list(args):
    """Look through all the types and releases and report which are built
    in a table to the user"""

    info = defaultdict(ReleaseInfo)

    # first, look for the type json files
    for f in Path(args.type_json).glob("*.yaml"):
        info[f.stem].has_type_json = True

    # Next, lets look at the package directory for releases
    for f in Path(args.type_package).glob("*"):
        info[f.name].has_package = True

    # Dump the info out as a rich table
    table = Table(title="Release Information")
    table.add_column("Release", justify="right", style="cyan", no_wrap=True)
    table.add_column("Type JSON", justify="right", style="magenta", no_wrap=True)
    table.add_column("Python Package", justify="right", style="green", no_wrap=True)

    for r in info:
        table.add_row(
            r,
            "Yes" if info[r].has_type_json else "No",
            "Yes" if info[r].has_package else "No",
        )

    console = Console()
    console.print(table)

def do_delete(args):
    """Delete the release packages"""
    for r in args.release:
        # Delete the type packages
        package_path = Path(args.type_package) / r
        if package_path.exists():
            logging.debug(f"Removing package {package_path}")
            shutil.rmtree(package_path)

        # Delete the type json file
        json_path = Path(args.type_json) / f"{r}.yaml"
        if json_path.exists():
            logging.debug(f"Removing type json {json_path}")
            json_path.unlink()
    return 0

def main():
    parser = argparse.ArgumentParser(
        prog="atlas_release_type_builder",
        description="Build and test atlas release types",
        epilog="Building type libraries can take considerable time. Unless the --clean"
        " option is used, previous outputs are used rather than regenerating from "
        "scratch.",
    )
    parser.add_argument(
        "-v",
        "--verbose",
        action="count",
        default=0,
        help="Verbosity Level (-v, or -vv))",
    )

    commands = parser.add_subparsers(help="sub-command help")

    # The "build" an atlas release command
    build_command = commands.add_parser(
        "build", help="Build type library for a release"
    )

    def add_build_args(parser, add_release=True, package_flags=True):
        if add_release:
            parser.add_argument(
                "release", type=str, help="List of releases to build", action="append"
            )
            parser.add_argument(
                "--clean",
                type=bool,
                action=argparse.BooleanOptionalAction,
                default=False,
            )
        parser.add_argument(
            "--type_json",
            type=Path,
            default=Path("../type_files"),
            help="Location where yaml type files should be written",
        )
        if package_flags:
            parser.add_argument(
                "--type_package",
                type=Path,
                default=Path(
                    "../type_packages",
                    help="Location where the python type package should be created",
                ),
            )
            parser.add_argument(
                "--package_version",
                type=str,
                default="2.0.0a1",
                help="Version of the python package to create",
            )
        parser.add_argument(
            "--command_location",
            type=Path,
            default=Path("../../"),
            help="Location of the scripts to run to build the type json file",
        )

    add_build_args(build_command)
    build_command.set_defaults(func=do_build)

    # The build yaml command - just does the first part of the build.
    build_yaml_command = commands.add_parser(
        "yaml", help="Build type yaml for a release"
    )
    add_build_args(build_yaml_command, package_flags=False)
    build_yaml_command.set_defaults(func=do_yaml)

    # The test command
    test_command = commands.add_parser("test", help="Run tests for a release")
    test_command.add_argument(
        "--test", choices=valid_tests, default=[], action="append"
    )
    test_command.add_argument(
        "--test_dir",
        type=Path,
        help="Path to place (and leave) test directory",
        default=None,
    )
    add_build_args(test_command)
    test_command.set_defaults(func=do_test)

    # The list command - list all releases that are built one way or the other.
    list_command = commands.add_parser("list", help="List all releases")
    list_command.set_defaults(func=do_list)
    add_build_args(list_command, add_release=False)

    # delete command - will delete a release(s) (type and package files).
    delete_command = commands.add_parser("delete", help="Delete a release")
    add_build_args(delete_command)
    delete_command.set_defaults(func=do_delete)

    args = parser.parse_args()

    # Global flags
    logging.basicConfig()
    if args.verbose:
        logging.getLogger().setLevel(logging.INFO)
        if int(args.verbose) > 1:
            logging.getLogger().setLevel(logging.DEBUG)

    # Now execute the command
    if not hasattr(args, "func"):
        logging.error("No command found on command line")
        parser.print_help()
        exit(1)

    sys.exit(args.func(args))

if __name__ == "__main__":
    main()


from pydantic import BaseModel

from service.util import remove_special_characters, strip_accents

class SummonerSchema(BaseModel):
    name: str
    level: int
    id: str

class ActiveGameSchema(BaseModel):
    gamemode: str
    participants: list

class ChampionSchema(BaseModel):
    champion_name: str
    title: str
    id: int
    image: str
    changes: dict

def parseSummoner(summoner: dict):
    return SummonerSchema(
        name=summoner["name"],
        level=summoner["summonerLevel"],
        id=summoner["id"]
    )

def parseActiveGame(active_game: dict):
    participants = _parseParticipants(active_game["participants"])
    return ActiveGameSchema(
        gamemode=active_game["gameMode"],
        participants=participants
    )

def _parseParticipants(participants: dict):
    blue_team = [
        [remove_special_characters(strip_accents(x["summonerName"])), x["championId"]]
        for x in participants
        if x["teamId"] == 100
    ]
    red_team = [
        [remove_special_characters(strip_accents(x["summonerName"])), x["championId"]]
        for x in participants
        if x["teamId"] == 200
    ]

    participants = blue_team + red_team
    parsed_participants = list()
    for participant in participants:
        parsed_participants.append({"summoner_name": participant[0], "champion_id": participant[1]})

    return parsed_participants


from torch.utils.data import Dataset, DataLoader
import os
import numpy as np
import re
import cv2
from pathlib import Path
import torchvision.transforms as transforms

def random_flip(image, steering_angle):
    """
    Randomly flipt the image left <-> right, and adjust the steering angle.
    """
    if np.random.rand() < 0.5:
        image = cv2.flip(image, 1)
        steering_angle = -steering_angle
    return image, steering_angle

def random_translate(image, steering_angle, range_x, range_y):
    """
    Randomly shift the image virtially and horizontally (translation).
    """
    trans_x = range_x * (np.random.rand() - 0.5)
    trans_y = range_y * (np.random.rand() - 0.5)
    steering_angle += trans_x * 0.002
    trans_m = np.float32([[1, 0, trans_x], [0, 1, trans_y]])
    height, width = image.shape[:2]
    image = cv2.warpAffine(image, trans_m, (width, height))
    return image, steering_angle

def augment(image, steering_angle, range_x=100, range_y=10):
    """
    Generate an augumented image and adjust steering angle.
    (The steering angle is associated with the center image)
    """
    image, steering_angle = random_flip(image, steering_angle)
    image, steering_angle = random_translate(image, steering_angle, range_x, range_y)
    return image, steering_angle

transformations = transforms.Compose([transforms.Lambda(lambda x: (x / 127.5) - 1.0)])

class DrivingData(Dataset):
    def __init__(
            self, file_path, device,
            target_size=(320, 240), # (width, height)
            crop_size = (200, 200)):

        self.processed_data_file = Path(file_path)
        self.device = device
        self.target_size = tuple(target_size)
        self.crop_size = tuple(crop_size)

        # Idea = associate each filename with a corresponding steering or label
        self.filenames = []
        self.imgs = None
        self.ground_truths = []

        if self.processed_data_file.is_file():
            self._load_processed_data()

        self.transform = transforms.Compose([
            transforms.ToTensor(),
            transforms.Lambda(lambda x: (x / 127.5) - 1.0)
        ])

    def _load_processed_data(self):
        if self.imgs is None:
            print("Loading processed data from {}...".format(self.processed_data_file))
            with np.load(self.processed_data_file) as data:
                self.imgs = data['imgs']
                self.ground_truths = data['ground_truths']
            print("Images: ", self.imgs.shape)
            print("Ground truths: ", self.ground_truths.shape)
            print("Done!")

    def _decode_experiment_dir(self, dir_subpath):
        # Load steerings or labels in the experiment dir
        steerings_filename = os.path.join(dir_subpath, "sync_steering.txt")
        labels_filename = os.path.join(dir_subpath, "labels.txt")

        # Try to load steerings first. Make sure that the steering angle or the
        # label file is in the first column. Note also that the first line are
        # comments so it should be skipped.
        try:
            ground_truths = np.loadtxt(steerings_filename, usecols=0,
                                  delimiter=',', skiprows=1)
            exp_type = 1
        except OSError as e:
            # Try load collision labels if there are no steerings
            try:
                ground_truths = np.loadtxt(labels_filename, usecols=0)
                exp_type = 0
            except OSError as e:
                print("Neither steerings nor labels found in dir {}".format(
                dir_subpath))
                raise IOError

        # Now fetch all images in the image subdir
        image_dir_path = os.path.join(dir_subpath, "images")
        for root, _, files in self._recursive_list(image_dir_path):
            sorted_files = sorted(files,
                    key = lambda fname: int(re.search(r'\d+',fname).group()))
            for frame_number, fname in enumerate(sorted_files):
                is_valid = False
                for extension in self.formats:
                    if fname.lower().endswith('.' + extension):
                        is_valid = True
                        break
                if is_valid:
                    absolute_path = os.path.join(root, fname)
                    self.filenames.append(os.path.relpath(absolute_path,
                            self.directory))
                    self.ground_truths.append(ground_truths[frame_number])
                    self.exp_type.append(exp_type)
                    self.samples += 1

    def _recursive_list(self, subpath):
        return sorted(os.walk(subpath, followlinks=False),
                key=lambda tpl: tpl[0])

    def __len__(self):
        return self.ground_truths.shape[0]

    def __getitem__(self, index):
        img = self.imgs[index]
        steering_angle = self.ground_truths[index]
        return img, steering_angle

def get_iterator_driving(file_path, device, batch_size=1, num_workers=0):
    """
    returns an iterator over DRIVING dataset batches
    :param file_path: path to .npz file containing a list of tuples
         each of them representing a path to an image and it class
    :param device:
    :param batch_size:
    :return: torch.utils.DataLoader object constructed from DRIVING dataset object
    """

    dataset = DrivingData(file_path, device=device)
    if 'test' in file_path or 'validation' in file_path:
        iterator = DataLoader(dataset, shuffle=False, batch_size=batch_size, num_workers=num_workers)
    else:
        iterator = DataLoader(dataset, shuffle=True, batch_size=batch_size, num_workers=num_workers)

    return iterator


# -*- coding: utf-8 -*-
from Action import Action
import util
import re
from datetime import datetime
import threading

class PrintSinaL2Action(Action):
    def __init__(self, **kwargs):
        Action.__init__(self, **kwargs )
        self._producer_list = [
            {
                 "name"             :    "SinaLevel2WS",
                 "producer_name"    :    "PrintSinaL2.SinaLevel2-quotation"
            }]
        self._init()
        self.logger.info(self._name +u"初始化")
        self.count = 0

    # 需要重写的方法
    def handler(self, event):
        dt = datetime.now()
        if isinstance(event.data, list):
            for data in event.data:
                self.count += 1
                print("PrintSinaL2:{}, \nCount:{},  Length: {}\n{}".format( dt,self.count, len(data),data ) )
        else:
            self.count += 1
            print("PrintSinaL2:{}, \nCount:{},  Length: {}\n{}".format( dt,self.count, len(event.data),event.data ) )

        # self.logger.info("线程退出")


# frozen_lake 학습해 성공률 계산해보기
# (아직 keras 잘 안되므로 김성훈교수님 2015 dqn 사용)keras로 frozen lake에 적용해보자.
# 5000 episode까지 학습한 성공률 0.6674.
# 학습 성공.
# 모델 저장 및 불러오기 추가해보자. (추가 완료)
# keras로도 바꿔보자.
import numpy as np
import tensorflow as tf
import random
from collections import deque
import dqn
import pylab
import gym
from typing import List

env = gym.make('FrozenLake-v0')
save_file = './model.ckpt'
INPUT_SIZE = env.observation_space.n    #one hot 인코딩해서 input size 16가능. one hot 인코딩 안하고 학습하니 학습이 아예 안된다!
OUTPUT_SIZE = env.action_space.n

DISCOUNT_RATE = 0.99
REPLAY_MEMORY = 50000
BATCH_SIZE = 64
TARGET_UPDATE_FREQUENCY = 5
MAX_EPISODES = 5000
LOAD_MODEL = True

def one_hot(x):
    return np.identity(16)[x:x+1]

def replay_train(mainDQN: dqn.DQN, targetDQN: dqn.DQN, train_batch: list) -> float:
    """Trains `mainDQN` with target Q values given by `targetDQN`
    Args:
        mainDQN (dqn.DQN): Main DQN that will be trained
        targetDQN (dqn.DQN): Target DQN that will predict Q_target
        train_batch (list): Minibatch of replay memory
            Each element is (s, a, r, s', done)
            [(state, action, reward, next_state, done), ...]
    Returns:
        float: After updating `mainDQN`, it returns a `loss`
    """
    states = np.vstack([x[0] for x in train_batch])
    actions = np.array([x[1] for x in train_batch])
    rewards = np.array([x[2] for x in train_batch])
    next_states = np.vstack([x[3] for x in train_batch])
    done = np.array([x[4] for x in train_batch])

    X = states

    Q_target = rewards + DISCOUNT_RATE * np.max(targetDQN.predict(next_states), axis=1) * ~done

    y = mainDQN.predict(states)
    y[np.arange(len(X)), actions] = Q_target

    # Train our network using target and predicted Q values on each episode
    return mainDQN.update(X, y)

def get_copy_var_ops(*, dest_scope_name: str, src_scope_name: str) -> List[tf.Operation]:
    """Creates TF operations that copy weights from `src_scope` to `dest_scope`
    Args:
        dest_scope_name (str): Destination weights (copy to)
        src_scope_name (str): Source weight (copy from)
    Returns:
        List[tf.Operation]: Update operations are created and returned
    """
    # Copy variables src_scope to dest_scope
    op_holder = []

    src_vars = tf.get_collection(
        tf.GraphKeys.TRAINABLE_VARIABLES, scope=src_scope_name)
    dest_vars = tf.get_collection(
        tf.GraphKeys.TRAINABLE_VARIABLES, scope=dest_scope_name)

    for src_var, dest_var in zip(src_vars, dest_vars):
        op_holder.append(dest_var.assign(src_var.value()))

    return op_holder

def bot_play(mainDQN: dqn.DQN, env: gym.Env) -> None:
    """Test runs with rendering and prints the total score
    Args:
        mainDQN (dqn.DQN): DQN agent to run a test
        env (gym.Env): Gym Environment
    """
    state = env.reset()
    reward_sum = 0

    while True:

        env.render()
        action = np.argmax(mainDQN.predict(state))
        state, reward, done, _ = env.step(action)
        reward_sum += reward

        if done:
            print("Total score: {}".format(reward_sum))
            break

def main():
    replay_buffer = deque(maxlen=REPLAY_MEMORY)

    with tf.Session() as sess:
        mainDQN = dqn.DQN(sess, INPUT_SIZE, OUTPUT_SIZE, name="main")
        targetDQN = dqn.DQN(sess, INPUT_SIZE, OUTPUT_SIZE, name="target")

        saver = tf.train.Saver()
        if LOAD_MODEL == False :
            sess.run(tf.global_variables_initializer())
        elif LOAD_MODEL == True :
            saver.restore(sess, save_file)

        # initial copy q_net -> target_net
        copy_ops = get_copy_var_ops(dest_scope_name="target",
                                    src_scope_name="main")
        sess.run(copy_ops)

        graph_x, graph_y = [],[]
        success_cnt = 0

        for episode in range(MAX_EPISODES):
            e = 1. / ((episode / 10) + 1)
            done = False
            step_count = 0
            state = env.reset()

            while not done:
                if np.random.rand() < e:
                    action = env.action_space.sample()
                else:
                    # Choose an action by greedily from the Q-network
                    action = np.argmax(mainDQN.predict(one_hot(state)))

                # Get new state and reward from environment
                next_state, reward, done, _ = env.step(action)

                if done:
                    if reward == 0: # 실패
                        reward = -10
                    else:   # 성공
                        reward = 10
                        success_cnt += 1

                # Save the experience to our buffer
                replay_buffer.append((one_hot(state), action, reward, one_hot(next_state), done))

                if len(replay_buffer) > BATCH_SIZE:
                    minibatch = random.sample(replay_buffer, BATCH_SIZE)
                    loss, _ = replay_train(mainDQN, targetDQN, minibatch)

                # 일정 주기로 타겟 업데이트
                if step_count % TARGET_UPDATE_FREQUENCY == 0:
                    sess.run(copy_ops)

                state = next_state
                step_count += 1

            graph_x.append(episode)
            graph_y.append(reward)
            if episode % 50 == 0 and episode != 0:
                pylab.plot(graph_x, graph_y)
                pylab.savefig("./frozen_save/frozen_lake_dqn.png")
                print("Episode: {} 까지 성공률: {}".format(episode, success_cnt / episode))
            if episode % 500 == 0 and episode != 0:
                saver.save(sess, save_file)  # 500주기마다 모델 save
                print("Episode: {} 모델 저장 완료".format(episode))

if __name__ == "__main__":
    main()


"""FastAPI で wasm (wasi) を実行して wasi の出力を返すサンプル."""
import json
import logging
from typing import Any

from aiofiles import tempfile
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from wasmtime import Engine, Linker, Module, Store, WasiConfig, WasmtimeError

app = FastAPI()

logging.basicConfig(level="INFO")
logger = logging.getLogger(__name__)

class Item(BaseModel):
    """Item."""

    name: str
    age: int
    description: str

@app.get("/{funcname}")
async def execute(funcname: str) -> Item:
    """Execute wasm function."""
    # refs: https://gist.github.com/pims/711549577759ad1341f1a90860f1f3a5

    wasm_path = f"functions/{funcname}.wasm"
    # create data
    d = {"name": "zztkm", "age": 20}

    try:
        item_dict = await invoke_wasm_module(funcname, wasm_path, d)
        return Item(**item_dict)
    except Exception:
        logger.exception("error")
        raise HTTPException(status_code=400, detail="Bad Request") from None

async def invoke_wasm_module(modname: str, wasm_path: str, input_data: dict[str, Any]) -> dict[str, Any]:
    """Invoke wasm module."""
    # refs: https://gist.github.com/pims/711549577759ad1341f1a90860f1f3a5

    engine = Engine()
    linker = Linker(engine)
    linker.define_wasi()

    # functions/funcname.wasm が存在するか確認
    try:
        module = Module.from_file(linker.engine, wasm_path)
    except Exception:
        logger.exception("module not found %s", wasm_path)
        raise

    config = WasiConfig()

    # generate tenmporary file for stdout_file
    async with tempfile.NamedTemporaryFile() as f:
        config.stdout_file = f.name
        config.env = [["DATA", json.dumps(input_data)]]

        store = Store(linker.engine)
        store.set_wasi(config)

        instance = linker.instantiate(store, module)
        start = instance.exports(store)["_start"]

        logger.info("start wasm module %s", modname)
        try:
            start(store)
            #await f.seek(0)
            out = await f.read()
            return json.loads(out.decode(encoding="utf-8"))
        except WasmtimeError as e:
            logger.debug(e)
            #await f.seek(0)
            out = await f.read()
            return json.loads(out.decode(encoding="utf-8"))


from Bio.Blast.Applications import NcbiblastnCommandline, NcbitblastnCommandline
from Bio.Blast import NCBIXML
from Bio import SeqIO
import subprocess

import csv, os, re, sys, glob

def BuildBlastDB(genome_folder, genome_fastafile, blastdatabasedir, BLASTbindir,buildblastdb):
	database_in = genome_folder+'/'+genome_fastafile
	genome = genome_fastafile.split('.fa')[0]
	database_store = blastdatabasedir + '/' + genome
	cmnd = BLASTbindir+'makeblastdb -dbtype nucl -in '+database_in+' -out '+database_store
	if buildblastdb == 'yes':
		print ("/ ---BUILDING BLAST DB FROM FASTA FILE---")
		print (cmnd)
		print((os.system(cmnd)))
	elif buildblastdb == 'no':
		print ("/ No BLAST db will be built\n\n")
		pass
	elif buildblastdb != 'yes':
		print ("// No BLAST db will be built, use set buildblastdb to 'yes' to build a BLAST database from fastafiles.")

	return database_store, genome

def BlastAndParse(blastdb, genome, queryfile, E_VALUE_THRESH, PERC_IDENTITY_THRESH, outputdir, BLAST_task, BLAST_type,output_contigs):
	blastoutput = outputdir+BLAST_type+queryfile.split('/')[-1]+"_vs"+genome+"_presence_absence.blastout"
	open(blastoutput, 'wb').close() #clearfile
	if BLAST_type == 'BLASTN':
		blastn_cline = NcbiblastnCommandline(query=queryfile, db=blastdb, evalue=E_VALUE_THRESH, outfmt=5, out=blastoutput, dust='no', task=BLAST_task)
		print("")
		print (blastn_cline)
		blastn_cline()
	else:
		print ("ERROR: NO BLAST TYPE HAS BEEN SPECIFIED (BLASTN OR TBLASTN)")

	result_handle = open(blastoutput)
	blast_records = NCBIXML.parse(result_handle)

	effector2scoresForThisGenome = {}

	for blast_record in blast_records:
		for alignment in blast_record.alignments:
			for hsp in alignment.hsps:
				return_perc_id = int(100*hsp.identities/blast_record.query_length)

				########
				return_contig_nr = "??"
				subject_name = alignment.title
				if "ontig_" in subject_name:
					return_contig_nr = subject_name.split("ontig_")[1]
					if "_" in return_contig_nr:
						return_contig_nr = return_contig_nr.split("_")[0]
				if blast_record.query not in effector2scoresForThisGenome: effector2scoresForThisGenome[blast_record.query]=[]
				if (hsp.expect < E_VALUE_THRESH) and (return_perc_id >= PERC_IDENTITY_THRESH):
					if output_contigs in 'no': 				#'', 'no' or 'n'
						effector2scoresForThisGenome[blast_record.query].append(return_perc_id)
					elif output_contigs in 'yes':
						effector2scoresForThisGenome[blast_record.query].append(return_contig_nr+'('+str(hsp.sbjct_start)+'-'+str(hsp.sbjct_end)+')')
	result_handle.close()
	return effector2scoresForThisGenome

class pres_abs_varApp():

	def __init__(self):
		self.verbose = False

	def start(self,force,show_all, queryfile, genome_folder, blastdatabasedir, outputdir, buildblastdb, r_location, BLASTbindir='', PERC_IDENTITY_THRESH=80, working=''):

	### arguments being passed from pipeline script: ###
		if working!='':
			if working[-1]!='/':
				working+='/'
			if '../' in working:
				dir = os.path.dirname(__file__)
				working = os.path.join(dir, working)
			os.chdir(working)
		else:
			working=os.path.dirname(__file__)
		print((os.getcwd()))

		if '../' in queryfile:
			dir = os.path.dirname(working)
			queryfile = os.path.join(dir, queryfile)
		if '../' in genome_folder:
			dir = os.path.dirname(working)
			genome_folder = os.path.join(dir, genome_folder)
			genome_folder = os.path.abspath(os.path.realpath(genome_folder))

		if '../' in blastdatabasedir:
			dir = os.path.dirname(working)
			blastdatabasedir = os.path.join(dir, blastdatabasedir)
			blastdatabasedir = os.path.abspath(os.path.realpath(blastdatabasedir))

		if BLASTbindir !='':
			if '../' in BLASTbindir or BLASTbindir[0]!='/':
				dir = os.path.dirname(working)
				BLASTbindir = os.path.join(dir, BLASTbindir)
				BLASTbindir = os.path.abspath(os.path.realpath(BLASTbindir))
			if BLASTbindir[:-1]!='/':
				BLASTbindir+='/'

		if '../' in outputdir:
			dir = os.path.dirname(working)
			outputdir = os.path.join(dir, outputdir)
			outputdir = os.path.abspath(os.path.realpath(outputdir))

		if '../' in r_location:
			dir = os.path.dirname(working)
			r_location = os.path.join(dir, r_location)

		if outputdir[-1]=="/":
			outputdir=outputdir[:-1]
		if str(genome_folder[-1])=="/":
			genome_folder=genome_folder[:-1]
		if queryfile[-1]=="/":
			queryfile=queryfile[:-1]

		queryfilename = queryfile.split('/')[-1]
		queryfolder = queryfile.split(queryfilename)[0]
		genome_folder_b = genome_folder.split('/')[-1]
		genome_folder_a = genome_folder.split(genome_folder_b)[0]
		BLAST_task="blastn"
		outputdir +='/03.blastn_presence_absence/'
		if not os.path.exists(outputdir):
			os.makedirs(outputdir)
		###
		#print buildblastdb, "test"
		#date_time_now 			= datetime.datetime.now().strftime("%d-%m-%Y %H:%M")
		outfilename 			= outputdir+genome_folder_b+'.vs.'+queryfilename.split('.fa')[0]+'.txt'
		E_VALUE_THRESH 			= 0.001
		BLAST_type 				= 'BLASTN'
		transpose 				= True
		plot_copynumber 		= False #define whether presence is '1' or copy nr (# of blast hits)

		if force !=True and os.path.isfile(outfilename)==True:
			x=2
			while os.path.isfile(outfilename)==True:
				outfilename = outputdir+genome_folder_b+'.vs.'+queryfilename.split('.fa')[0]+'_run_'+str(x)+'.txt'
				x+=1

		if buildblastdb == 'yes':
			print ("// Will start with building BlastDBs of each genome encountered.\n")

		yes = set(['yes','y'])
		no = set(['no','n', ''])
		output_contigs = 'n'
		hierarchicalclusteringquestion = 'y'
		if hierarchicalclusteringquestion in yes:
			hierarchicalclustering = True
			outfilename = outfilename.split('.txt')[0]+'_hierarch-clust.txt'

		else:
			hierarchicalclustering = False
			output_contigs = eval(input("Would you like to get contig numbers in stead of %id? (y/n) \n>"))
			if output_contigs in yes:
				outfilename = outfilename.split('.txt')[0]+'_contignrs.txt'

		if transpose == True:
			outfilename = outfilename.split('.txt')[0]+'_transposed.txt'
		query_list = []

		for record in SeqIO.parse(queryfile, 'fasta'):
			query_list.append(record.id)
	   	###############################################################
	   	#cancels all above outputfile naming statements:
		outfilename = outputdir+'blastn_presence_absence.txt'
	  	###############################################################
		effector2genome2scores = {}
		genomes = []
		for genome_fastafile in os.listdir(genome_folder):
			if genome_fastafile.endswith((".fa", ".fasta", ".fas")):
				blastdb, genome = BuildBlastDB(genome_folder, genome_fastafile, blastdatabasedir, BLASTbindir,buildblastdb)
				genomes.append(genome)
				#return a dictionary containing all query fastaheaders (keys) with BLAST% as values:
				effector2scoresForThisGenome = BlastAndParse(blastdb, genome, queryfile, E_VALUE_THRESH, PERC_IDENTITY_THRESH, outputdir, BLAST_task,BLAST_type,output_contigs)

				#order these back to the way they were in the original query fastafile:
				n=1
				for effector in query_list:
					new_record_id = '{0:04}'.format(n)+'..'+effector 			#'..' to reduce mis-splitting
					if new_record_id not in effector2genome2scores: effector2genome2scores[new_record_id] = {}
					#in dictionary effector2scoresForThisGenome, add empty values (for effectors that were not found with BLAST:
					if effector not in effector2scoresForThisGenome:
						effector2scoresForThisGenome[effector] = []
					effector2genome2scores[new_record_id][genome] = effector2scoresForThisGenome[effector]
					n+=1
			else:
				print(('-'*20))
				print(("// No more files with extension .fa, .fas, .fasta were found in directory '%s'" % (genome_folder)))
				print(('-'*20))

		# Print to table:

		open(outfilename, 'w').close() 									#clear file
		outfile     = open(outfilename, 'a')

		genomes.sort()
		header = ''
		if hierarchicalclustering == False:
			if output_contigs in no:
				header = BLAST_type+' PERCENT IDENTITY\nE-value: '+str(E_VALUE_THRESH)+'\n'+'%ID_threshold: '+str(PERC_IDENTITY_THRESH)+'%\n''query:'
			elif output_contigs in yes:
				header = BLAST_type+' CONTIG NR + SUBJECT LOCATION\nE-value: '+str(E_VALUE_THRESH)+'\n'+'%ID_threshold: '+str(PERC_IDENTITY_THRESH)+'%\n''query:'

		if transpose == True:
			for effector in sorted(effector2genome2scores.keys()):
				header += '\t'+effector
		else:
			for g in genomes:
				header += '\t'+g
		header+='\n'
		print (header)
		outfile.write(header)

		if transpose == True:
			for g in genomes:
				out = g
				for effector in sorted(effector2genome2scores.keys()):
					genome2scores = effector2genome2scores[effector]
					scores = genome2scores[g]
					if hierarchicalclustering == True:
						scorestr = 0
						for s in scores:
							if plot_copynumber == True:
								scorestr	+=1
							else:
								scorestr	=1
					else:
						scorestr = ''
						for s in scores:
							scorestr+=str(s)+';'
						scorestr = scorestr[:-1] # laatste ; eraf
					out += '\t'+str(scorestr)
				print (out)
				outfile.write(out+'\n')
		else:
			for effector in sorted(effector2genome2scores.keys()):
				out = effector
				genome2scores = effector2genome2scores[effector]
				for g in genomes:
					scores = genome2scores[g]
					if hierarchicalclustering == True:
						scorestr = 0
						for s in scores:
							scorestr=1
					else:
						scorestr = ''
						for s in scores:
							scorestr+=str(s)+';'
						scorestr = scorestr[:-1] # laatste ; eraf
					out += '\t'+str(scorestr)
				print (out)
				outfile.write(out+'\n')
		outfile.close()
		print(('-'*30))
		print(("// Written data to file: ", outfilename))
		print(('-'*30,'\n', "Generating image"))
		print((r_location, outfilename))

		if show_all==False:
			subprocess.call ("Rscript --vanilla "+str(r_location)+" "+str(outfilename)+" FALSE", shell=True)
		else:
			subprocess.call ("Rscript --vanilla "+str(r_location)+" "+str(outfilename)+" TRUE", shell=True)


import scipy.io
import os
import numpy as np
import cv2
from src.dataset_loader.datasetUtils import *
from src.module.function import *

# AEI_kmeans = np.array(
#     [[-0.37959412,  0.0896221,   0.05287072,  0.90319805,  0.99185765,  0.99559656],
#      [-0.83103055,  0.10626562,  0.0299934 ,  0.51789511,  0.9870735 ,  0.99758095],
#      [ 0.96297726,  0.68287062,  0.09528844, -0.21831189,  0.72485644,  0.99080621],
#      [ 0.53015532,  0.10014143,  0.02805288, -0.83132568,  0.99192195,  0.99745991],
#      [ 0.86861489,  0.10909378, -0.01550593,  0.45770347,  0.98740085,  0.99779991],
#      [-0.93655641,  0.08003212,  0.00114281, -0.29706842,  0.98999588,  0.99767647],
#      [ 0.02219816,  0.03370241,  0.00365841, -0.98923496,  0.9974772 ,  0.99869783],
#      [ 0.92599029,  0.07402331,  0.01549187, -0.33420168,  0.99187946,  0.99774758],
#      [-0.52019853,  0.10331311, -0.03536063, -0.8353981 ,  0.98985149,  0.99392136],
#      [ 0.45297786,  0.10298805, -0.06113805,  0.8753308 ,  0.98987731,  0.99512509]]
# )

AEI_kmeans = np.array(
    [
     [-0.47621441,  0.12350617,  0.0664845  , 0.87744119 , 0.9890067  , 0.9958762 ],
     [-0.79701111,  0.0897479 ,  0.01711335 , 0.59947227 , 0.99189595 , 0.99781593],
     [ 0.24718982,  0.03466764, -0.04019557 , 0.96675895 , 0.99782538 , 0.99749265],
     [-0.92900045,  0.12796586, -0.02959894 ,-0.36313565 , 0.98995881 , 0.998478],
     [-0.99093385,  0.30963827,  0.04399168 ,-0.03849395 , 0.94132017 , 0.99682202],
     [ 0.63991822,  0.08676728,  0.01266549, -0.76207316 , 0.99382095,  0.99825609],
     [-0.97041808, -0.04349611,  0.05179524, -0.2217277 ,  0.99625264,  0.99575635],
     [-0.75931431,  0.34058989,  0.08872104,  0.64092271 , 0.93278592,  0.99194674],
     [ 0.36390551,  0.10402673,  0.04212833, -0.92760847 , 0.99118484,  0.9969696 ],
     [-0.02068825,  0.04477927,  0.00249808,  0.99825918 , 0.99533083,  0.99785909],
     [-0.93060167,  0.09966637,  0.02097796,  0.35872446 , 0.99066755,  0.99873714],
     [ 0.559026  ,  0.1759681 , -0.07896671,  0.82671968,  0.98174557,  0.99525479],
     [-0.98538106, -0.04002435,  0.01965351,  0.13296903,  0.99573206,  0.99765323],
     [ 0.98356559,  0.00115694, -0.00883092, -0.14943953,  0.99380634,  0.99798697],
     [-0.63086791,  0.09089955,  0.04917653,  0.77313566,  0.99302347,  0.9975129 ],
     [ 0.96297726,  0.68287062,  0.09528844 ,-0.21831189,  0.72485644,  0.99080621],
     [ 0.89273093,  0.11288946,  0.02832545, -0.43880685,  0.99064152 , 0.99760898],
     [-0.36066788,  0.23448267,  0.3728395 ,  0.92279967,  0.95020208 , 0.91946087],
     [ 0.37875795,  0.19846965, -0.14137089,  0.92281866,  0.97781774 , 0.98612405],
     [ 0.43960992,  0.05180215 ,-0.03794313,  0.89593427,  0.99690567,  0.99759366],
     [ 0.95803985,  0.08486974 ,-0.00948285,  0.27001179,  0.98916025,  0.99752139],
     [-0.3050887 ,  0.06243963 , 0.05041205,  0.95028344,  0.99524137,  0.99639665],
     [-0.74702989,  0.11269824 ,-0.01525002, -0.65815991,  0.99057759,  0.9981728 ],
     [-0.29650311,  0.05262828 ,-0.0096012 , -0.95147564,  0.99610075,  0.99613081],
     [-0.51872658,  0.27368644 ,-0.43666126, -0.8388962 ,  0.94688394,  0.89374456],
     [-0.5108259 ,  0.11407817 ,-0.03612973, -0.85536679,  0.98887529,  0.99624544],
     [ 0.68169805,  0.07487039 ,-0.02669289,  0.72808247,  0.99466217,  0.99823521],
     [ 0.84553889,  0.14025878 ,-0.02059248,  0.52720822,  0.98423765,  0.99785371],
     [ 0.65659539,  0.52835238 ,-0.2444696,  0.74462741 , 0.84188589,  0.94337074],
     [ 0.01983669,  0.02214991,  0.00109092, -0.99691492 , 0.99825084,  0.99899499]
    ]
)

class dataLoaderSingleObject(object):
    def __init__(self,
                 Pascal3DDataPath=None,
                 trainOrVal='train',
                 ):
        self.dataStart = 0
        self.dataLength = 0
        self.epoch = 0

        self._Pascal3DDataPath = Pascal3DDataPath
        self._trainOrVal = trainOrVal
        self._isTrain = True
        if self._trainOrVal == 'train':
            self._isTrain = True
        elif self._trainOrVal == 'val':
            self._isTrain = False
        else:
            print('choose \'train\' or \'val\'')
            return
        self._dataPathList = []
        self._CAD3DShapes = {}

        print('set pascal3d dataset...')

        self._getTrainList()
        self._loadDataPath()
        self._load3DShapes()
        self._dataPathShuffle()

    def _getTrainList(self):
        print('set train or val list...')
        self._Pascal3DTrainList = []
        datasetList = os.listdir(os.path.join(self._Pascal3DDataPath, 'Image_sets/'))
        for datasetName in datasetList:
            if os.path.isdir(os.path.join(self._Pascal3DDataPath, 'Image_sets/', datasetName)):
                txtFileList = os.listdir(os.path.join(self._Pascal3DDataPath, 'Image_sets/',datasetName))
                for txtFileName in txtFileList:
                    className = txtFileName.split('.')[0].split('_')[0]
                    trainval = txtFileName.split('.')[0].split('_')[-1]
                    if trainval==self._trainOrVal:
                        with open(os.path.join(self._Pascal3DDataPath, 'Image_sets/',datasetName,txtFileName)) as txtFilePointer:
                            dataPointList = txtFilePointer.readlines()
                            for i, dataPoint in enumerate(dataPointList):
                                if datasetName == 'pascal':
                                    dp = dataPoint.split('\n')[0].split(' ')[0]
                                    isTrue = int(dataPoint.split('\n')[0].split(' ')[-1])
                                    if int(isTrue)==1:
                                        self._Pascal3DTrainList.append(dp)
                                else:
                                    dp = dataPoint.split('\n')[0].split(' ')[0]
                                    self._Pascal3DTrainList.append(dp)
        print('done!')

    def _loadDataPath(self):
        print('load datapoint path...')
        datasetList = os.listdir(os.path.join(self._Pascal3DDataPath, 'training_data'))
        for datasetName in datasetList:
            if datasetName == 'imagenet' or datasetName == 'pascal':
                dataPointList = os.listdir(os.path.join(self._Pascal3DDataPath, 'training_data', datasetName))
                for dataPointName in dataPointList:
                    if dataPointName in self._Pascal3DTrainList:
                        if os.path.isdir(os.path.join(self._Pascal3DDataPath, 'training_data', datasetName, dataPointName)):
                            dataPointPath = os.path.join(self._Pascal3DDataPath, 'training_data', datasetName, dataPointName)
                            self._dataPathList.append(dataPointPath)
        self._dataPathList = np.array(self._dataPathList)
        self.dataLength = len(self._dataPathList)
        print('done!')

    def _dataPathShuffle(self):
        # print('')
        # print('data path shuffle...')
        self.dataStart = 0
        np.random.shuffle(self._dataPathList)
        self.dataLength = len(self._dataPathList)
        # print('done! : ' + str(self.dataLength))

    def _load3DShapes(self):
        print('load 3d shapes for pascal3d...')
        self._classIndex = {}
        self._CAD3DShapes = {}
        classList = os.listdir(os.path.join(self._Pascal3DDataPath, 'CAD'))
        classList.sort(key=natural_keys)
        classNum = 0
        for className in classList:
            if os.path.isdir(os.path.join(self._Pascal3DDataPath, 'CAD', className)):
                if self._CAD3DShapes.get(className) is None:
                    self._classIndex[className] = classNum
                    self._CAD3DShapes[className] = []
                    classNum += 1
                CADModelList = os.listdir(os.path.join(self._Pascal3DDataPath, 'CAD', className))
                CADModelList.sort(key=natural_keys)
                for CADModel in CADModelList:
                    if CADModel.split(".")[-1] == 'npy':
                        shape = np.load(os.path.join(self._Pascal3DDataPath, 'CAD', className, CADModel)).reshape((64, 64, 64, 1))
                        self._CAD3DShapes[className].append(shape)
                self._CAD3DShapes[className] = np.array(self._CAD3DShapes[className])
                self._CAD3DShapes[className] = np.where(self._CAD3DShapes[className] > 0, 1.0, 0.0)
        self._classNum = classNum
        print('done!')

    def getKmeansAEI(self, k=None, max_iter=1000):
        AEI = []
        i = 0
        for dataPath in self._dataPathList:
            print(i,len(self._dataPathList))
            i += 1
            objFolderList = os.listdir(dataPath)
            objSelectedList = []
            for objFolder in objFolderList:
                if os.path.isdir(os.path.join(dataPath, objFolder)):
                    objInfoTXT = os.path.join(dataPath, objFolder, 'objInfo.txt')
                    with open(objInfoTXT) as objInfoPointer:
                        objInfo = objInfoPointer.readline()
                    className = objInfo.split(' ')[0]
                    if self._CAD3DShapes[className] is not None:
                        objSelectedList.append(objInfo)
                    for objIndex, objSelected in enumerate(objSelectedList):
                        className,imagePath,CADModelPath,colMin,rowMin,colMax,rowMax,azimuth,elevation,inPlaneRot=objSelected.split(' ')
                        azimuth, elevation, inPlaneRot = float(azimuth) / 180.0 * np.pi, float(elevation) / 180.0 * np.pi, float(inPlaneRot) / 180.0 * np.pi
                        sin = np.sin([azimuth, elevation, inPlaneRot])
                        cos = np.cos([azimuth, elevation, inPlaneRot])
                        AEI.append([sin[0], sin[1], sin[2], cos[0], cos[1], cos[2]])
        AEI = np.array(AEI)
        kmeans = kmeans_cosine(X=AEI, k=k, max_iter=max_iter)
        centers, xtoc, dist = kmeans.kmeansSample()
        return centers, xtoc, dist

    def getNextBatch(self, batchSizeof3DShape=32, imageSize=None, augmentation=True):
        inputImages = []
        outputImages, classList, instList, EulerRadList = [], [], [], []
        item_num = 0
        while len(outputImages)==0:
            for dataPath in self._dataPathList[self.dataStart:]:
                objFolderList = os.listdir(dataPath)
                np.random.shuffle(objFolderList)
                # objFolderList.sort()
                objSelectedList = []
                for objFolder in objFolderList:
                    objFolder = objFolder.decode("utf-8")
                    if os.path.isdir(os.path.join(dataPath, objFolder)):
                        objInfoTXT = os.path.join(dataPath, objFolder, 'objInfo.txt')
                        with open(objInfoTXT) as objInfoPointer:
                            objInfo = objInfoPointer.readline()
                        className = objInfo.split(' ')[0]
                        # if className == 'car':
                        #     objSelectedList.append(objInfo)
                        if self._CAD3DShapes[className] is not None:
                            objSelectedList.append(objInfo)
                np.random.shuffle(objSelectedList)
                if len(objSelectedList) > 0:
                    if len(outputImages) >= batchSizeof3DShape and len(inputImages)>0:
                        break
                    try:
                        image2DPath = os.path.join(self._Pascal3DDataPath, objSelectedList[0].split(' ')[1])
                        inputImage = cv2.imread(image2DPath, cv2.IMREAD_COLOR)
                        row_org, col_org, _ = inputImage.shape
                        for objIndex, objSelected in enumerate(objSelectedList):
                            className,imagePath,CADModelPath,colMin,rowMin,colMax,rowMax,azimuth,elevation,inPlaneRot=objSelected.split(' ')
                            colMin, rowMin, colMax, rowMax, azimuth, elevation, inPlaneRot = float(colMin), float(rowMin), float(
                                colMax), float(rowMax), float(azimuth), float(elevation), float(inPlaneRot)
                            width, height = colMax - colMin, rowMax - rowMin
                            if augmentation:
                                border_ratio = np.random.rand() * 0.2
                            else:
                                border_ratio = 0.1
                            colMin, rowMin = np.max((0, colMin - width * border_ratio)), np.max(
                                (0, rowMin - height * border_ratio))
                            colMax, rowMax = np.min((col_org, colMax + width * border_ratio)), np.min(
                                (row_org, rowMax + height * border_ratio))
                            image = inputImage[int(rowMin):int(rowMax), int(colMin):int(colMax), :]

                            is_flip = False
                            if augmentation:
                                is_flip = np.random.rand() > 0.5
                            if is_flip:
                                image = cv2.flip(image, 1)
                                if np.random.rand() > 0.5:
                                    image = cv2.flip(image, 0)
                            # image augmentation
                            image, _, _, _, _, _ = imageRandomAugmentation(
                                inputImage=image, imageRowFinal=imageSize[0], imageColFinal=imageSize[1],
                                imAug=augmentation, imAugPr=0.5, padding=False,
                                randomTrans=augmentation, randomScale=augmentation,
                                transPr=0.5, scalePr=0.5, transRatioMax=0.2,
                                scaleRatioMin=0.8, scaleRatioMax=1.2
                            )
                            image = image / 255.

                            azimuth, elevation, inPlaneRot = float(azimuth) / 180.0 * np.pi, float(elevation) / 180.0 * np.pi, float(inPlaneRot) / 180.0 * np.pi
                            # for crossdataset
                            # inPlaneRot = 0.
                            if is_flip:
                                azimuth = -azimuth
                            EulerRad = np.array([azimuth, elevation, inPlaneRot])
                            classVector = np.zeros(self._classNum)
                            classVector[self._classIndex[className]] = 1
                            cadIndex = int(CADModelPath.split('/')[-1])
                            instVector = np.zeros(10)
                            instVector[cadIndex - 1] = 1
                            # car 3d shape
                            if (cadIndex - 1 < 0) or (cadIndex > len(self._CAD3DShapes)):
                                for i in range(1000):
                                    print('pascal ', cadIndex - 1)
                                    return
                            inst3DCAD = self._CAD3DShapes[className][cadIndex - 1]

                            inputImages.append(image)
                            outputImages.append(inst3DCAD)
                            classList.append(classVector)
                            instList.append(instVector)
                            EulerRadList.append(EulerRad)
                            item_num += 1
                            if item_num >= batchSizeof3DShape:
                                break
                    except:
                        pass
                self.dataStart += 1
                if self.dataStart >= self.dataLength:
                    self.epoch += 1
                    self._dataPathShuffle()
                    break
        inputImages = np.array(inputImages).astype('float32')
        outputImages = np.array(outputImages).astype('float32')
        classList = np.array(classList).astype('float32')
        instList = np.array(instList).astype('float32')
        EulerRadList = np.array(EulerRadList).astype('float32')

        return instList, classList, np.sin(EulerRadList), np.cos(EulerRadList), inputImages, outputImages

# average image size of pascal3D : (508.54, 404.47)
class dataLoader(object):
    def __init__(self,
                 imageSize=(640,480),
                 gridSize=(20,15),
                 predNumPerGrid=5,
                 Pascal3DDataPath=None,
                 trainOrVal='train',
                 ):
        self.dataStart = 0
        self.dataLength = 0
        self.epoch = 0

        self._imageSize = imageSize
        self._gridSize = gridSize
        self._predNumPerGrid = predNumPerGrid
        self._Pascal3DDataPath = Pascal3DDataPath
        self._trainOrVal = trainOrVal
        self._isTrain = True
        if self._trainOrVal == 'train':
            self._isTrain = True
        elif self._trainOrVal == 'val':
            self._isTrain = False
        else:
            print('choose \'train\' or \'val\'')
            return
        self._dataPathList = []
        self._CAD3DShapes = None

        print('set pascal3d dataset...')

        self._getTrainList()
        self._loadDataPath()
        self._load3DShapes()
        self._dataPathShuffle()

    def _getTrainList(self):
        print('set train or val list...')
        self._Pascal3DTrainList = []
        datasetList = os.listdir(os.path.join(self._Pascal3DDataPath, 'Image_sets/'))
        for datasetName in datasetList:
            if os.path.isdir(os.path.join(self._Pascal3DDataPath, 'Image_sets/', datasetName)):
                txtFileList = os.listdir(os.path.join(self._Pascal3DDataPath, 'Image_sets/',datasetName))
                for txtFileName in txtFileList:
                    className = txtFileName.split('.')[0].split('_')[0]
                    trainval = txtFileName.split('.')[0].split('_')[-1]
                    if className=='car' and trainval==self._trainOrVal:
                        with open(os.path.join(self._Pascal3DDataPath, 'Image_sets/',datasetName,txtFileName)) as txtFilePointer:
                            dataPointList = txtFilePointer.readlines()
                            for i, dataPoint in enumerate(dataPointList):
                                if datasetName == 'pascal':
                                    dp = dataPoint.split('\n')[0].split(' ')[0]
                                    isTrue = int(dataPoint.split('\n')[0].split(' ')[-1])
                                    if int(isTrue)==1:
                                        self._Pascal3DTrainList.append(dp)
                                else:
                                    dp = dataPoint.split('\n')[0].split(' ')[0]
                                    self._Pascal3DTrainList.append(dp)
        print('done!')

    def _loadDataPath(self):
        print('load datapoint path...')
        datasetList = os.listdir(os.path.join(self._Pascal3DDataPath, 'training_data'))
        for datasetName in datasetList:
            if datasetName == 'imagenet' or datasetName == 'pascal':
                dataPointList = os.listdir(os.path.join(self._Pascal3DDataPath, 'training_data', datasetName))
                for dataPointName in dataPointList:
                    if dataPointName in self._Pascal3DTrainList:
                        if os.path.isdir(os.path.join(self._Pascal3DDataPath, 'training_data', datasetName, dataPointName)):
                            dataPointPath = os.path.join(self._Pascal3DDataPath, 'training_data', datasetName, dataPointName)
                            self._dataPathList.append(dataPointPath)
        self._dataPathList = np.array(self._dataPathList)
        self.dataLength = len(self._dataPathList)
        print('done!')

    def _dataPathShuffle(self):
        print('')
        print('data path shuffle...')
        self.dataStart = 0
        np.random.shuffle(self._dataPathList)
        self.dataLength = len(self._dataPathList)
        print('done! : ' + str(self.dataLength))

    def _load3DShapes(self):
        print('load 3d shapes for pascal3d...')
        self._CAD3DShapes = []
        CADModelList = os.listdir(os.path.join(self._Pascal3DDataPath, 'CAD', 'car'))
        CADModelList.sort()
        for CADModel in CADModelList:
            if CADModel.split(".")[-1] == 'npy':
                shape = np.load(os.path.join(self._Pascal3DDataPath, 'CAD', 'car', CADModel)).reshape(64, 64, 64, 1)
                self._CAD3DShapes.append(shape)
        self._CAD3DShapes = np.array(self._CAD3DShapes)
        self._CAD3DShapes = np.where(self._CAD3DShapes>0, 1.0, 0.0)
        print('done!')

    def _getOffset(self, batchSize):
        offsetX = np.transpose(np.reshape(
            np.array([np.arange(self._gridSize[0])]*self._gridSize[1]*self._predNumPerGrid),
            (self._predNumPerGrid, self._gridSize[1], self._gridSize[0])), (1,2,0))
        offsetX = np.tile(np.reshape(offsetX, (1,self._gridSize[1],self._gridSize[0],self._predNumPerGrid)),[batchSize,1,1,1])
        offsetY = np.transpose(np.reshape(
            np.array([np.arange(self._gridSize[1])]*self._gridSize[0]*self._predNumPerGrid),
            (self._predNumPerGrid, self._gridSize[0], self._gridSize[1])), (2,1,0))
        offsetY = np.tile(np.reshape(offsetY, (1,self._gridSize[1],self._gridSize[0],self._predNumPerGrid)),[batchSize,1,1,1])
        return offsetX.astype('float32'), offsetY.astype('float32')

    def getNextBatch(self, batchSizeof3DShape=32, imageSize=None, gridSize=None):
        if imageSize!=None:
            self._imageSize = imageSize
        if gridSize!=None:
            self._gridSize = gridSize
        inputImages, bboxImages, objnessImages, eulerRadImages = [], [], [], []
        outputImages, instList = [], []
        while len(outputImages)==0:
            for dataPath in self._dataPathList[self.dataStart:]:
                objFolderList = os.listdir(dataPath)
                np.random.shuffle(objFolderList)
                # objFolderList.sort()
                objSelectedList = []
                for objFolder in objFolderList:
                    if os.path.isdir(os.path.join(dataPath, objFolder)):
                        objInfoTXT = os.path.join(dataPath, objFolder, 'objInfo.txt')
                        with open(objInfoTXT) as objInfoPointer:
                            objInfo = objInfoPointer.readline()
                        className = objInfo.split(' ')[0]
                        if className == 'car':
                            objSelectedList.append(objInfo)
                if len(objSelectedList) > 0:
                    if len(outputImages) + len(objSelectedList) > batchSizeof3DShape and len(inputImages)>0:
                        break
                    try:
                        image2DPath = os.path.join(self._Pascal3DDataPath, objSelectedList[0].split(' ')[1])
                        inputImage = cv2.imread(image2DPath, cv2.IMREAD_COLOR)

                        is_flip = np.random.rand() > 0.5
                        if is_flip:
                            inputImage = cv2.flip(inputImage, flipCode=1)

                        # imageRow, imageCol, channel = inputImage.shape
                        # # make this a square image
                        # imgRowColMax = np.max((imageRow, imageCol))
                        # heightBorderSize = (imgRowColMax - imageRow) / 2
                        # widthBorderSize = (imgRowColMax - imageCol) / 2
                        # inputImage = cv2.copyMakeBorder(
                        #     inputImage, top=heightBorderSize, bottom=heightBorderSize,
                        #     left=widthBorderSize, right=widthBorderSize, borderType=cv2.BORDER_CONSTANT,
                        #     value=[0, 0, 0])

                        imageRowOrg, imageColOrg, _ = inputImage.shape
                        inputImage, rowPad, dRow, colPad, dCol, scale = imageRandomAugmentation(
                            inputImage=inputImage, imageRowFinal=self._imageSize[1], imageColFinal=self._imageSize[0],
                            imAug=self._isTrain, imAugPr=0.5, padding=False,
                            randomTrans=self._isTrain, randomScale=self._isTrain,
                            transPr=0.5, scalePr=0.5, transRatioMax=0.2,
                            scaleRatioMin=0.8, scaleRatioMax=1.2
                        )
                        inputImage = inputImage/255.

                        imageRow = imageRowOrg + 2*rowPad
                        imageCol = imageColOrg + 2*colPad

                        objOrderingImage = -1 * np.ones([self._gridSize[1], self._gridSize[0], self._predNumPerGrid])
                        bboxImage = np.zeros([self._gridSize[1], self._gridSize[0], self._predNumPerGrid, 4])
                        objnessImage = np.zeros([self._gridSize[1], self._gridSize[0], self._predNumPerGrid, 1])
                        eulerRadImage = np.zeros([self._gridSize[1], self._gridSize[0], self._predNumPerGrid, 3])
                        outputPerImage, instPerImage, EulerPerImage = [], [], []
                        itemIndex = 0
                        for objIndex, objSelected in enumerate(objSelectedList):
                            className,imagePath,CADModelPath,colMin,rowMin,colMax,rowMax,azimuth,elevation,inPlaneRot=objSelected.split(' ')
                            colMin, rowMin, colMax, rowMax, azimuth, elevation = float(colMin), float(rowMin), float(
                                colMax), float(rowMax), float(azimuth), float(elevation)
                            if is_flip:
                                colMin, colMax = imageColOrg - colMax, imageColOrg - colMin
                                azimuth, inPlaneRot = -azimuth, -inPlaneRot
                            # #border for square
                            # rowMin, rowMax = float(rowMin) + heightBorderSize, float(rowMax) + heightBorderSize
                            # colMin, colMax = float(colMin) + widthBorderSize, float(colMax) + widthBorderSize
                            # augmentation
                            colMin = (float(colMin) + colPad - imageCol/2) * scale + imageCol/2 + dCol
                            colMax = (float(colMax) + colPad - imageCol/2) * scale + imageCol/2 + dCol
                            rowMin = (float(rowMin) + rowPad - imageRow/2) * scale + imageRow/2 + dRow
                            rowMax = (float(rowMax) + rowPad - imageRow/2) * scale + imageRow/2 + dRow
                            # colMin = np.max((0.0, colMin))
                            # colMax = np.min((colMax, imageCol))
                            # rowMin = np.max((0.0, rowMin))
                            # rowMax = np.min((rowMax, imageRow))
                            azimuth,elevation,inPlaneRot = -float(azimuth)/180.0*np.pi,-float(elevation)/180.0*np.pi,float(inPlaneRot)/180.0*np.pi
                            cadIndex = int(CADModelPath.split('/')[-1])

                            rowCenterOnGrid = (rowMax+rowMin)/2.0*self._gridSize[1]/imageRow
                            colCenterOnGrid = (colMax+colMin)/2.0*self._gridSize[0]/imageCol
                            rowIndexOnGrid = int(rowCenterOnGrid)
                            colIndexOnGrid = int(colCenterOnGrid)
                            dx,dy = colCenterOnGrid - colIndexOnGrid, rowCenterOnGrid - rowIndexOnGrid
                            bboxHeight,bboxWidth = np.min((1.0, (rowMax-rowMin)/imageRow)),np.min((1.0, (colMax-colMin)/imageCol))
                            for predIndex in range(self._predNumPerGrid):
                                if (rowIndexOnGrid>=0 and rowIndexOnGrid<self._gridSize[1]) \
                                    and (colIndexOnGrid>=0 and colIndexOnGrid<self._gridSize[0]) \
                                    and bboxHeight > 0 and bboxWidth > 0 \
                                    and objnessImage[rowIndexOnGrid, colIndexOnGrid, predIndex] != 1:
                                    # objectness and bounding box
                                    objnessImage[rowIndexOnGrid,colIndexOnGrid,predIndex]=1
                                    bboxImage[rowIndexOnGrid,colIndexOnGrid,predIndex,0:4] = bboxHeight,bboxWidth,dx,dy
                                    eulerRadImage[rowIndexOnGrid,colIndexOnGrid,predIndex,:] = azimuth, elevation, inPlaneRot

                                    # car instance vector
                                    carInstVector = np.zeros(len(self._CAD3DShapes))
                                    carInstVector[cadIndex-1] = 1
                                    # car 3d shape
                                    if(cadIndex-1<0) or (cadIndex>len(self._CAD3DShapes)):
                                        for i in range(1000):
                                            print('pascal ', cadIndex-1)
                                            return
                                    car3DCAD = self._CAD3DShapes[cadIndex-1]
                                    # # Euler angle in rad
                                    # EulerRad = np.array([azimuth,elevation,inPlaneRot])

                                    # append items
                                    outputPerImage.append(car3DCAD)
                                    instPerImage.append(carInstVector)
                                    # EulerPerImage.append(EulerRad)
                                    # set item order
                                    objOrderingImage[rowIndexOnGrid, colIndexOnGrid, predIndex] = itemIndex
                                    itemIndex += 1
                                    break
                        if itemIndex > 0:
                            inputImages.append(inputImage)
                            bboxImages.append(bboxImage)
                            objnessImages.append(objnessImage)
                            eulerRadImages.append(eulerRadImage)

                            for gridRow in range(self._gridSize[1]):
                                for gridCol in range(self._gridSize[0]):
                                    for predIndex in range(self._predNumPerGrid):
                                        objOrder = int(objOrderingImage[gridRow, gridCol, predIndex])
                                        if objOrder>=0:
                                            outputImages.append(outputPerImage[objOrder])
                                            instList.append(instPerImage[objOrder])
                                            # EulerRadList.append(EulerPerImage[objOrder])
                    except:
                        pass
                self.dataStart += 1
                if self.dataStart >= self.dataLength:
                    self.epoch += 1
                    self._dataPathShuffle()
                    break
        inputImages = np.array(inputImages).astype('float32')
        objnessImages = np.array(objnessImages).astype('float32')
        bboxImages = np.array(bboxImages).astype('float32')
        eulerRadImages = np.array(eulerRadImages).astype('float32')
        outputImages = np.array(outputImages).astype('float32')
        instList = np.array(instList).astype('float32')
        # EulerRadList = np.array(EulerRadList).astype('float32')

        offsetX,offsetY = self._getOffset(batchSize=len(inputImages))

        # print inputImages.shape
        # print instList.shape
    #     return offsetX, offsetY, inputImages, objnessImages,\
    # bboxImages, np.sin(EulerRadList), np.cos(EulerRadList), \
    # outputImages, instList
        return offsetX, offsetY, inputImages, objnessImages, \
               bboxImages, np.sin(eulerRadImages), np.cos(eulerRadImages), \
               outputImages, instList

# data_loader_pascal = dataLoaderSingleObject(trainOrVal='val',
#                                              Pascal3DDataPath='/media/yonsei/4TB_HDD/dataset/PASCAL3D+_release1.1/')
#
# center, xtoc, dist = data_loader_pascal.getKmeansAEI(k=30*12)
# print(center)
# print(np.median(dist))
# np.save('../../val.npy', center)