"""
setup earthquake depth relocation directory
"""

import obspy
import sh
import numpy as np
import click
from os.path import join
from glob import glob
import copy


def generate_new_cmtsolution_files(cmts_dir, generated_cmts_dir, depth_perturbation_list):
    cmt_names = glob(join(cmts_dir, "*"))
    for cmt_file in cmt_names:
        event = obspy.read_events(cmt_file)[0]
        # gcmt_id = event.resource_id.id.split("/")[-2]
        # there are some problems in changing names
        gcmt_id = cmt_file.split("/")[-1]

        # assume dirs like f"{generated_cmts_dir}/d-3" have already been created
        for depth_per in depth_perturbation_list:
            generated_name = join(generated_cmts_dir, f"d{depth_per}", gcmt_id)
            # there are always problem in copy event, so here I'd like to read in the event again
            event_this_depth = obspy.read_events(cmt_file)[0]
            # event_this_depth = event.copy()
            event_this_depth.origins[0].depth += 1000.0*depth_per
            # print(generated_name, generated_cmts_dir, f"d{depth_per}", gcmt_id)
            event_this_depth.write(generated_name, format="CMTSOLUTION")


def setup_basic_structure(main_dir, ref_dir, cmts_dir, depth_perturbation_list):
    # main
    sh.mkdir("-p", main_dir)

    # ref
    sh.cp("-r", ref_dir, join(main_dir, "ref"))

    # refine the structure in ref
    sh.rm("-rf", join(main_dir, "ref", "DATABASES_MPI"))
    sh.rm("-rf", join(main_dir, "ref", "EXAMPLES"))
    sh.rm("-rf", join(main_dir, "ref", "OUTPUT_FILES"))
    sh.rm("-rf", join(main_dir, "ref", "doc"))
    sh.rm("-rf", join(main_dir, "ref", "tests"))

    # mv DATA and utils to upper level
    sh.mv(join(main_dir, "ref", "DATA"), main_dir)
    sh.mv(join(main_dir, "ref", "utils"), main_dir)

    # cmts
    sh.mkdir("-p", join(main_dir, "cmts"))
    sh.cp("-r", cmts_dir, join(main_dir, "cmts", "cmts_raw"))
    sh.mkdir("-p", join(main_dir, "cmts", "cmts_generated"))
    for depth_per in depth_perturbation_list:
        sh.mkdir("-p", join(main_dir, "cmts",
                            "cmts_generated", f"d{depth_per}"))

    # working directory
    sh.mkdir("-p", join(main_dir, "work"))


def setup_structure_after_generat_cmts(main_dir, output_dir, depth_perturbation_list):
    # get cmts names
    cmt_dirs = glob(join(main_dir, "cmts", "cmts_raw", "*"))
    cmt_names = [item.split("/")[-1] for item in cmt_dirs]

    # mkdirs
    for cmt_name in cmt_names:
        sh.mkdir(join(main_dir, "work", cmt_name))
        for depth_per in depth_perturbation_list:
            # sh.mkdir(join(main_dir, "work", cmt_name, f"d{depth_per}"))
            # cp ref to working dirs
            sh.cp("-r", join(main_dir, "ref"),
                  join(main_dir, "work", cmt_name, f"d{depth_per}"))

    # mv DATA and utils back to ref
    sh.mv(join(main_dir, "DATA"), join(main_dir, "ref", "DATA"))
    sh.mv(join(main_dir, "utils"), join(main_dir, "ref", "utils"))

    # mkdir DATA in work directory
    for cmt_name in cmt_names:
        for depth_per in depth_perturbation_list:
            sh.mkdir(join(main_dir, "work", cmt_name, f"d{depth_per}", "DATA"))

    # cp and ln files in DATA
    toln = ["cemRequest", "crust1.0", "crust2.0",
            "crustmap", "epcrust", "eucrust-07", "GLL", "heterogen", "Lebedev_sea99", "Montagner_model", "old", "PPM", "QRFSI12", "s20rts", "s362ani", "s40rts", "Simons_model", "topo_bathy", "Zhao_JP_model"]
    for cmt_name in cmt_names:
        for depth_per in depth_perturbation_list:
            sh.cp(join(main_dir, "cmts", "cmts_generated",
                       f"d{depth_per}", cmt_name), join(main_dir, "work", cmt_name, f"d{depth_per}", "DATA", "CMTSOLUTION"))
            sh.cp(join(main_dir, "ref", "DATA", "Par_file"), join(
                main_dir, "work", cmt_name, f"d{depth_per}", "DATA", "Par_file"))
            sh.cp(join(main_dir, "ref", "DATA", "STATIONS"), join(
                main_dir, "work", cmt_name, f"d{depth_per}", "DATA", "STATIONS"))
            for lnfile in toln:
                sh.ln("-s", join(main_dir, "ref", "DATA", lnfile), join(
                    main_dir, "work", cmt_name, f"d{depth_per}", "DATA", lnfile))

            # ln in work files
            toln_work = ["utils"]
            for lnfile in toln_work:
                sh.ln("-s", join(main_dir, "ref", lnfile), join(
                    main_dir, "work", cmt_name, f"d{depth_per}", lnfile))

    # mkdir and ln DATABASE_MPI and OUTPUT_FILES
    sh.mkdir("-p", output_dir)
    sh.mkdir("-p", join(output_dir, "DATABASES_MPI"))
    sh.mkdir("-p", join(output_dir, "OUTPUT_FILES"))
    for cmt_name in cmt_names:
        for depth_per in depth_perturbation_list:
            sh.mkdir("-p", join(output_dir, "DATABASES_MPI",
                                cmt_name, f"d{depth_per}"))
            sh.mkdir("-p", join(output_dir, "OUTPUT_FILES",
                                cmt_name, f"d{depth_per}"))
            sh.ln("-s", join(output_dir, "DATABASES_MPI",
                             cmt_name, f"d{depth_per}"), join(main_dir, "work", cmt_name, f"d{depth_per}", "DATABASES_MPI"))
            sh.ln("-s", join(output_dir, "OUTPUT_FILES",
                             cmt_name, f"d{depth_per}"), join(main_dir, "work", cmt_name, f"d{depth_per}", "OUTPUT_FILES"))


@click.command()
@click.option('--main_dir', required=True, help="the main working directory", type=str)
@click.option('--output_dir', required=True, help="the output directory in scratch", type=str)
@click.option('--ref_dir', required=True, help="the reference specfem directory", type=str)
@click.option('--cmts_dir', required=True, help="the cmt solution directory", type=str)
@click.option('--depth_perturbation', required=True, help="the depth perturbation, use somthing like -3,-1,5 (in km)", type=str)
def main(main_dir, output_dir, ref_dir, cmts_dir, depth_perturbation):
    depth_perturbation_list = [float(item)
                               for item in depth_perturbation.split(",")]
    setup_basic_structure(main_dir, ref_dir, cmts_dir, depth_perturbation_list)
    generated_cmts_dir = join(main_dir, "cmts", "cmts_generated")
    working_cmts_dir = join(main_dir, "cmts", "cmts_raw")
    generate_new_cmtsolution_files(
        working_cmts_dir, generated_cmts_dir, depth_perturbation_list)
    setup_structure_after_generat_cmts(
        main_dir, output_dir, depth_perturbation_list)


if __name__ == "__main__":
    main()

[((5319, 5334), 'click.command', 'click.command', ([], {}), '()\n', (5332, 5334), False, 'import click\n'), ((5336, 5426), 'click.option', 'click.option', (['"""--main_dir"""'], {'required': '(True)', 'help': '"""the main working directory"""', 'type': 'str'}), "('--main_dir', required=True, help='the main working directory',\n    type=str)\n", (5348, 5426), False, 'import click\n'), ((5424, 5522), 'click.option', 'click.option', (['"""--output_dir"""'], {'required': '(True)', 'help': '"""the output directory in scratch"""', 'type': 'str'}), "('--output_dir', required=True, help=\n    'the output directory in scratch', type=str)\n", (5436, 5522), False, 'import click\n'), ((5519, 5614), 'click.option', 'click.option', (['"""--ref_dir"""'], {'required': '(True)', 'help': '"""the reference specfem directory"""', 'type': 'str'}), "('--ref_dir', required=True, help=\n    'the reference specfem directory', type=str)\n", (5531, 5614), False, 'import click\n'), ((5611, 5701), 'click.option', 'click.option', (['"""--cmts_dir"""'], {'required': '(True)', 'help': '"""the cmt solution directory"""', 'type': 'str'}), "('--cmts_dir', required=True, help='the cmt solution directory',\n    type=str)\n", (5623, 5701), False, 'import click\n'), ((5699, 5831), 'click.option', 'click.option', (['"""--depth_perturbation"""'], {'required': '(True)', 'help': '"""the depth perturbation, use somthing like -3,-1,5 (in km)"""', 'type': 'str'}), "('--depth_perturbation', required=True, help=\n    'the depth perturbation, use somthing like -3,-1,5 (in km)', type=str)\n", (5711, 5831), False, 'import click\n'), ((1267, 1291), 'sh.mkdir', 'sh.mkdir', (['"""-p"""', 'main_dir'], {}), "('-p', main_dir)\n", (1275, 1291), False, 'import sh\n'), ((4496, 4522), 'sh.mkdir', 'sh.mkdir', (['"""-p"""', 'output_dir'], {}), "('-p', output_dir)\n", (4504, 4522), False, 'import sh\n'), ((6120, 6160), 'os.path.join', 'join', (['main_dir', '"""cmts"""', '"""cmts_generated"""'], {}), "(main_dir, 'cmts', 'cmts_generated')\n", (6124, 6160), False, 'from os.path import join\n'), ((6184, 6218), 'os.path.join', 'join', (['main_dir', '"""cmts"""', '"""cmts_raw"""'], {}), "(main_dir, 'cmts', 'cmts_raw')\n", (6188, 6218), False, 'from os.path import join\n'), ((281, 300), 'os.path.join', 'join', (['cmts_dir', '"""*"""'], {}), "(cmts_dir, '*')\n", (285, 300), False, 'from os.path import join\n'), ((1328, 1349), 'os.path.join', 'join', (['main_dir', '"""ref"""'], {}), "(main_dir, 'ref')\n", (1332, 1349), False, 'from os.path import join\n'), ((1403, 1441), 'os.path.join', 'join', (['main_dir', '"""ref"""', '"""DATABASES_MPI"""'], {}), "(main_dir, 'ref', 'DATABASES_MPI')\n", (1407, 1441), False, 'from os.path import join\n'), ((1460, 1493), 'os.path.join', 'join', (['main_dir', '"""ref"""', '"""EXAMPLES"""'], {}), "(main_dir, 'ref', 'EXAMPLES')\n", (1464, 1493), False, 'from os.path import join\n'), ((1512, 1549), 'os.path.join', 'join', (['main_dir', '"""ref"""', '"""OUTPUT_FILES"""'], {}), "(main_dir, 'ref', 'OUTPUT_FILES')\n", (1516, 1549), False, 'from os.path import join\n'), ((1568, 1596), 'os.path.join', 'join', (['main_dir', '"""ref"""', '"""doc"""'], {}), "(main_dir, 'ref', 'doc')\n", (1572, 1596), False, 'from os.path import join\n'), ((1615, 1645), 'os.path.join', 'join', (['main_dir', '"""ref"""', '"""tests"""'], {}), "(main_dir, 'ref', 'tests')\n", (1619, 1645), False, 'from os.path import join\n'), ((1697, 1726), 'os.path.join', 'join', (['main_dir', '"""ref"""', '"""DATA"""'], {}), "(main_dir, 'ref', 'DATA')\n", (1701, 1726), False, 'from os.path import join\n'), ((1748, 1778), 'os.path.join', 'join', (['main_dir', '"""ref"""', '"""utils"""'], {}), "(main_dir, 'ref', 'utils')\n", (1752, 1778), False, 'from os.path import join\n'), ((1821, 1843), 'os.path.join', 'join', (['main_dir', '"""cmts"""'], {}), "(main_dir, 'cmts')\n", (1825, 1843), False, 'from os.path import join\n'), ((1871, 1905), 'os.path.join', 'join', (['main_dir', '"""cmts"""', '"""cmts_raw"""'], {}), "(main_dir, 'cmts', 'cmts_raw')\n", (1875, 1905), False, 'from os.path import join\n'), ((1926, 1966), 'os.path.join', 'join', (['main_dir', '"""cmts"""', '"""cmts_generated"""'], {}), "(main_dir, 'cmts', 'cmts_generated')\n", (1930, 1966), False, 'from os.path import join\n'), ((2168, 2190), 'os.path.join', 'join', (['main_dir', '"""work"""'], {}), "(main_dir, 'work')\n", (2172, 2190), False, 'from os.path import join\n'), ((2322, 2361), 'os.path.join', 'join', (['main_dir', '"""cmts"""', '"""cmts_raw"""', '"""*"""'], {}), "(main_dir, 'cmts', 'cmts_raw', '*')\n", (2326, 2361), False, 'from os.path import join\n'), ((2842, 2864), 'os.path.join', 'join', (['main_dir', '"""DATA"""'], {}), "(main_dir, 'DATA')\n", (2846, 2864), False, 'from os.path import join\n'), ((2866, 2895), 'os.path.join', 'join', (['main_dir', '"""ref"""', '"""DATA"""'], {}), "(main_dir, 'ref', 'DATA')\n", (2870, 2895), False, 'from os.path import join\n'), ((2907, 2930), 'os.path.join', 'join', (['main_dir', '"""utils"""'], {}), "(main_dir, 'utils')\n", (2911, 2930), False, 'from os.path import join\n'), ((2932, 2962), 'os.path.join', 'join', (['main_dir', '"""ref"""', '"""utils"""'], {}), "(main_dir, 'ref', 'utils')\n", (2936, 2962), False, 'from os.path import join\n'), ((4542, 4575), 'os.path.join', 'join', (['output_dir', '"""DATABASES_MPI"""'], {}), "(output_dir, 'DATABASES_MPI')\n", (4546, 4575), False, 'from os.path import join\n'), ((4596, 4628), 'os.path.join', 'join', (['output_dir', '"""OUTPUT_FILES"""'], {}), "(output_dir, 'OUTPUT_FILES')\n", (4600, 4628), False, 'from os.path import join\n'), ((349, 376), 'obspy.read_events', 'obspy.read_events', (['cmt_file'], {}), '(cmt_file)\n', (366, 376), False, 'import obspy\n'), ((691, 741), 'os.path.join', 'join', (['generated_cmts_dir', 'f"""d{depth_per}"""', 'gcmt_id'], {}), "(generated_cmts_dir, f'd{depth_per}', gcmt_id)\n", (695, 741), False, 'from os.path import join\n'), ((2037, 2094), 'os.path.join', 'join', (['main_dir', '"""cmts"""', '"""cmts_generated"""', 'f"""d{depth_per}"""'], {}), "(main_dir, 'cmts', 'cmts_generated', f'd{depth_per}')\n", (2041, 2094), False, 'from os.path import join\n'), ((2484, 2516), 'os.path.join', 'join', (['main_dir', '"""work"""', 'cmt_name'], {}), "(main_dir, 'work', cmt_name)\n", (2488, 2516), False, 'from os.path import join\n'), ((871, 898), 'obspy.read_events', 'obspy.read_events', (['cmt_file'], {}), '(cmt_file)\n', (888, 898), False, 'import obspy\n'), ((2703, 2724), 'os.path.join', 'join', (['main_dir', '"""ref"""'], {}), "(main_dir, 'ref')\n", (2707, 2724), False, 'from os.path import join\n'), ((2744, 2793), 'os.path.join', 'join', (['main_dir', '"""work"""', 'cmt_name', 'f"""d{depth_per}"""'], {}), "(main_dir, 'work', cmt_name, f'd{depth_per}')\n", (2748, 2793), False, 'from os.path import join\n'), ((3102, 3159), 'os.path.join', 'join', (['main_dir', '"""work"""', 'cmt_name', 'f"""d{depth_per}"""', '"""DATA"""'], {}), "(main_dir, 'work', cmt_name, f'd{depth_per}', 'DATA')\n", (3106, 3159), False, 'from os.path import join\n'), ((3549, 3616), 'os.path.join', 'join', (['main_dir', '"""cmts"""', '"""cmts_generated"""', 'f"""d{depth_per}"""', 'cmt_name'], {}), "(main_dir, 'cmts', 'cmts_generated', f'd{depth_per}', cmt_name)\n", (3553, 3616), False, 'from os.path import join\n'), ((3641, 3713), 'os.path.join', 'join', (['main_dir', '"""work"""', 'cmt_name', 'f"""d{depth_per}"""', '"""DATA"""', '"""CMTSOLUTION"""'], {}), "(main_dir, 'work', cmt_name, f'd{depth_per}', 'DATA', 'CMTSOLUTION')\n", (3645, 3713), False, 'from os.path import join\n'), ((3733, 3774), 'os.path.join', 'join', (['main_dir', '"""ref"""', '"""DATA"""', '"""Par_file"""'], {}), "(main_dir, 'ref', 'DATA', 'Par_file')\n", (3737, 3774), False, 'from os.path import join\n'), ((3776, 3845), 'os.path.join', 'join', (['main_dir', '"""work"""', 'cmt_name', 'f"""d{depth_per}"""', '"""DATA"""', '"""Par_file"""'], {}), "(main_dir, 'work', cmt_name, f'd{depth_per}', 'DATA', 'Par_file')\n", (3780, 3845), False, 'from os.path import join\n'), ((3882, 3923), 'os.path.join', 'join', (['main_dir', '"""ref"""', '"""DATA"""', '"""STATIONS"""'], {}), "(main_dir, 'ref', 'DATA', 'STATIONS')\n", (3886, 3923), False, 'from os.path import join\n'), ((3925, 3994), 'os.path.join', 'join', (['main_dir', '"""work"""', 'cmt_name', 'f"""d{depth_per}"""', '"""DATA"""', '"""STATIONS"""'], {}), "(main_dir, 'work', cmt_name, f'd{depth_per}', 'DATA', 'STATIONS')\n", (3929, 3994), False, 'from os.path import join\n'), ((4738, 4798), 'os.path.join', 'join', (['output_dir', '"""DATABASES_MPI"""', 'cmt_name', 'f"""d{depth_per}"""'], {}), "(output_dir, 'DATABASES_MPI', cmt_name, f'd{depth_per}')\n", (4742, 4798), False, 'from os.path import join\n'), ((4859, 4918), 'os.path.join', 'join', (['output_dir', '"""OUTPUT_FILES"""', 'cmt_name', 'f"""d{depth_per}"""'], {}), "(output_dir, 'OUTPUT_FILES', cmt_name, f'd{depth_per}')\n", (4863, 4918), False, 'from os.path import join\n'), ((4976, 5036), 'os.path.join', 'join', (['output_dir', '"""DATABASES_MPI"""', 'cmt_name', 'f"""d{depth_per}"""'], {}), "(output_dir, 'DATABASES_MPI', cmt_name, f'd{depth_per}')\n", (4980, 5036), False, 'from os.path import join\n'), ((5067, 5133), 'os.path.join', 'join', (['main_dir', '"""work"""', 'cmt_name', 'f"""d{depth_per}"""', '"""DATABASES_MPI"""'], {}), "(main_dir, 'work', cmt_name, f'd{depth_per}', 'DATABASES_MPI')\n", (5071, 5133), False, 'from os.path import join\n'), ((5159, 5218), 'os.path.join', 'join', (['output_dir', '"""OUTPUT_FILES"""', 'cmt_name', 'f"""d{depth_per}"""'], {}), "(output_dir, 'OUTPUT_FILES', cmt_name, f'd{depth_per}')\n", (5163, 5218), False, 'from os.path import join\n'), ((5249, 5314), 'os.path.join', 'join', (['main_dir', '"""work"""', 'cmt_name', 'f"""d{depth_per}"""', '"""OUTPUT_FILES"""'], {}), "(main_dir, 'work', cmt_name, f'd{depth_per}', 'OUTPUT_FILES')\n", (5253, 5314), False, 'from os.path import join\n'), ((4073, 4110), 'os.path.join', 'join', (['main_dir', '"""ref"""', '"""DATA"""', 'lnfile'], {}), "(main_dir, 'ref', 'DATA', lnfile)\n", (4077, 4110), False, 'from os.path import join\n'), ((4112, 4177), 'os.path.join', 'join', (['main_dir', '"""work"""', 'cmt_name', 'f"""d{depth_per}"""', '"""DATA"""', 'lnfile'], {}), "(main_dir, 'work', cmt_name, f'd{depth_per}', 'DATA', lnfile)\n", (4116, 4177), False, 'from os.path import join\n'), ((4331, 4360), 'os.path.join', 'join', (['main_dir', '"""ref"""', 'lnfile'], {}), "(main_dir, 'ref', lnfile)\n", (4335, 4360), False, 'from os.path import join\n'), ((4362, 4419), 'os.path.join', 'join', (['main_dir', '"""work"""', 'cmt_name', 'f"""d{depth_per}"""', 'lnfile'], {}), "(main_dir, 'work', cmt_name, f'd{depth_per}', lnfile)\n", (4366, 4419), False, 'from os.path import join\n')]