#!/bin/python3
# Create an "experiment package" from "scene asset" and "experiment data".
# Package dir structure:
#  /
#  |-scene
#    |-meta.json
#    |-scene
#    | |-...
#    |-EqSound-0.wav
#    |- ...
#    |-collision-0.av
#    |- ...
import sys
sys.path.append('sound')
import argparse
import logging
import os
import os.path
import json
import shutil
import tarfile
import math
import numpy as np
import data_to_sound
import simulate_collision


logging.basicConfig(level=logging.DEBUG)
logger = logging.getLogger(__name__)


def is_scene_asset_shallow(path):
    """
    Perform a shallow check to see if the path
    is indeed scene asset generated by recon.
    """
    return (
        os.path.isdir(path) and
        os.path.isfile(os.path.join(path, "meta.json")))


class InvalidSceneAsset(Exception):
    def __init__(self, msg):
        self.msg = msg

    def __str__(self):
        return self.msg


class InvalidExpDesc(Exception):
    def __init__(self, msg):
        self.msg = msg

    def __str__(self):
        return self.msg


def copy_stripped_scene(src, dst):
    """
    Copy scene asset from src to dst, removing debug and checkpoint
    information if they exist.
    throws InvalidSceneAsset when encountering unexpected files etc.
    """
    def ignore(src, names):
        prefix_whitelist = [
            "sm_", "mat_", "diffuse_",
            "collision_", "meta.json", ]
        if os.path.abspath(src) == os.path.abspath(src):
            black = []
            for name in names:
                white = any(name.startswith(pw) for pw in prefix_whitelist)
                if not white:
                    black.append(name)
            return set(black)
        else:
            return set()
    logger.info("Copying stripped scene asset to %s", dst)
    copytree_wo_permission(src, dst, ignore)


def copytree_wo_permission(src, dst, ignore=None):
    """
    Recursively copy a directory tree, without trying to
    change permissions.
    """
    names = os.listdir(src)
    if ignore is not None:
        ignored_names = ignore(src, names)
    else:
        ignored_names = set()

    os.makedirs(dst)
    errors = []
    for name in names:
        if name in ignored_names:
            continue
        srcname = os.path.join(src, name)
        dstname = os.path.join(dst, name)
        try:
            if os.path.isdir(srcname):
                copytree_wo_permission(srcname, dstname, ignore)
            else:
                shutil.copyfile(srcname, dstname)
        # catch the Error from the recursive copytree so that we can
        # continue with other files
        except Error as err:
            errors.extend(err.args[0])
        except EnvironmentError as why:
            errors.append((srcname, dstname, str(why)))
    if errors:
        raise Error(errors)


def parse_knet(fobj):
    """
    Parse K-NET ascii file.
    http://www.kyoshin.bosai.go.jp/kyoshin/man/knetform.html

    return: {
        "accel": acceleration [in cm/s^2]
        "freq": sampling frequency
    }
    """
    metadata = {}
    for i in range(17):
        entry = fobj.readline()
        key = entry[:18].rstrip().decode('ascii')
        value = entry[18:].rstrip().decode('utf8')
        metadata[key] = value

    if metadata["Scale Factor"] != "3920(gal)/6182761":
        raise Exception("Unknown scale factor in k-net file")

    values = []
    for line in fobj:
        values += list(map(float, line.split()))
    acc = np.array(values) * (3920 / 6182761)
    if metadata["Sampling Freq(Hz)"] != "100Hz":
        raise Exception("Unknown sampling frequency")

    return {
        "accel": acc,
        "freq": 100
    }


def process_quake(package_path, knet_path, quake_desc):
    """
    Write asset files and return meta info
    Return json
    """
    if quake_desc["type"] != "K-NET":
        raise InvalidExpDesc("Unknown quake type: %s" % quake_desc["type"])
    # Open K-Net archive
    archive_filename = "%s%s.tar.gz" % (quake_desc["point"], quake_desc["date"])
    archive_path = os.path.join(knet_path, archive_filename)
    archive = tarfile.open(archive_path, "r")
    logger.info("Opening K-NET archive file at %s", archive_path)

    acc_data = []
    for channel in ["EW", "NS", "UD"]:
        knet_fobj = archive.extractfile(
            "./%s%s.%s" % (quake_desc["point"], quake_desc["date"], channel))
        acc_data.append(parse_knet(knet_fobj)["accel"])
    acc3d = np.array(acc_data).T
    freq = 100
    logger.info("Loaded %s-shaped acceleration at %f Hz" % (acc3d.shape, freq))

    # Choose subregion
    ts = np.arange(len(acc3d), dtype=float) / freq
    acc3d = acc3d[(quake_desc["t_begin"] <= ts) & (ts <= quake_desc["t_end"])]
    logger.info("Choosing %d samples of interest", len(acc3d))
    logger.info("Max gal: %f", np.max(np.abs(acc3d)))

    # Output as sound.
    bg_name = "bg-%s.wav" % quake_desc["point"]
    bg_asset = "bg-%s" % quake_desc["point"]
    bg_path = os.path.join(package_path, bg_name)
    data_to_sound.convert_acc_to_sound(acc3d, freq, bg_path)

    # Output as acceleration wave.
    framerate = 60
    assert(framerate < freq)
    low_acc3d = []
    for axis in range(3):
        samples = simulate_collision.butter_lowpass_filter(
            acc3d[:, axis], framerate / 2, freq)
        low_acc3d.append(samples)
    low_acc3d = np.array(low_acc3d).T
    logger.info("Max gal (after LPF): %f", np.max(np.abs(low_acc3d)))

    return {
        "bg_sound": bg_name,
        "bg_sound:asset": bg_asset,
        "accel": {
            "acc": [list(map(float, v)) for v in low_acc3d],
            "freq": freq
        }
    }


if __name__ == '__main__':
    parser = argparse.ArgumentParser(
        description="""
Generate experiment package from the following:
1. Scene Asset (directory), generated by recon
2. experiment config (json file)

Result will be generated in another directory.
""",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument(
        '--scene-asset', required=True, type=str,
        help='Directory path containing scene asset.')
    parser.add_argument(
        '--experiment', required=True, type=str,
        help='Path to experiment setup json file.')
    parser.add_argument(
        '--k-net', required=True, type=str,
        help="Path containing K-NET tar.gz files with names like 'MYG0041103111446.tar.gz'")
    parser.add_argument(
        '--output', required=True, type=str,
        help='Output wave file path for soundscape simulation.')

    args = parser.parse_args()

    if not is_scene_asset_shallow(args.scene_asset):
        logger.warn("%s doesn't seem to be a scene asset" % args.scene_asset)
        sys.exit(1)

    experiment = json.load(open(args.experiment))

    # Create package
    package_path = args.output
    if os.path.isfile(package_path):
        logger.warn("package path %s already exists, and it's not a package; aborting", package_path)
        sys.exit(1)
    elif os.path.isdir(package_path):
        logger.info("Deleting existing directory %s", package_path)
        shutil.rmtree(package_path)
    logger.info("Creating experiment package at %s", package_path)
    os.mkdir(package_path)

    package_meta = {}
    package_meta["quakes"] = []
    copy_stripped_scene(args.scene_asset, os.path.join(package_path, "scene"))
    for quake_desc in experiment["quakes"]:
        package_meta["quakes"].append(
            process_quake(package_path, args.k_net, quake_desc))

    json.dump(package_meta, open(os.path.join(package_path, "meta.json"), "w"))
