"""
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()

[((123, 1, 123, 16), 'click.command', 'click.command', ({}, {}), '()', False, 'import click\n'), ((124, 1, 124, 87), 'click.option', 'click.option', (), '', False, 'import click\n'), ((125, 1, 125, 94), 'click.option', 'click.option', (), '', False, 'import click\n'), ((126, 1, 126, 91), 'click.option', 'click.option', (), '', False, 'import click\n'), ((127, 1, 127, 87), 'click.option', 'click.option', (), '', False, 'import click\n'), ((128, 1, 128, 128), 'click.option', 'click.option', (), '', False, 'import click\n'), ((35, 4, 35, 28), 'sh.mkdir', 'sh.mkdir', ({(35, 13, 35, 17): '"""-p"""', (35, 19, 35, 27): 'main_dir'}, {}), "('-p', main_dir)", False, 'import sh\n'), ((108, 4, 108, 30), 'sh.mkdir', 'sh.mkdir', ({(108, 13, 108, 17): '"""-p"""', (108, 19, 108, 29): 'output_dir'}, {}), "('-p', output_dir)", False, 'import sh\n'), ((133, 25, 133, 65), 'os.path.join', 'join', ({(133, 30, 133, 38): 'main_dir', (133, 40, 133, 46): '"""cmts"""', (133, 48, 133, 64): '"""cmts_generated"""'}, {}), "(main_dir, 'cmts', 'cmts_generated')", False, 'from os.path import join\n'), ((134, 23, 134, 57), 'os.path.join', 'join', ({(134, 28, 134, 36): 'main_dir', (134, 38, 134, 44): '"""cmts"""', (134, 46, 134, 56): '"""cmts_raw"""'}, {}), "(main_dir, 'cmts', 'cmts_raw')", False, 'from os.path import join\n'), ((15, 21, 15, 40), 'os.path.join', 'join', ({(15, 26, 15, 34): 'cmts_dir', (15, 36, 15, 39): '"""*"""'}, {}), "(cmts_dir, '*')", False, 'from os.path import join\n'), ((38, 25, 38, 46), 'os.path.join', 'join', ({(38, 30, 38, 38): 'main_dir', (38, 40, 38, 45): '"""ref"""'}, {}), "(main_dir, 'ref')", False, 'from os.path import join\n'), ((41, 17, 41, 55), 'os.path.join', 'join', ({(41, 22, 41, 30): 'main_dir', (41, 32, 41, 37): '"""ref"""', (41, 39, 41, 54): '"""DATABASES_MPI"""'}, {}), "(main_dir, 'ref', 'DATABASES_MPI')", False, 'from os.path import join\n'), ((42, 17, 42, 50), 'os.path.join', 'join', ({(42, 22, 42, 30): 'main_dir', (42, 32, 42, 37): '"""ref"""', (42, 39, 42, 49): '"""EXAMPLES"""'}, {}), "(main_dir, 'ref', 'EXAMPLES')", False, 'from os.path import join\n'), ((43, 17, 43, 54), 'os.path.join', 'join', ({(43, 22, 43, 30): 'main_dir', (43, 32, 43, 37): '"""ref"""', (43, 39, 43, 53): '"""OUTPUT_FILES"""'}, {}), "(main_dir, 'ref', 'OUTPUT_FILES')", False, 'from os.path import join\n'), ((44, 17, 44, 45), 'os.path.join', 'join', ({(44, 22, 44, 30): 'main_dir', (44, 32, 44, 37): '"""ref"""', (44, 39, 44, 44): '"""doc"""'}, {}), "(main_dir, 'ref', 'doc')", False, 'from os.path import join\n'), ((45, 17, 45, 47), 'os.path.join', 'join', ({(45, 22, 45, 30): 'main_dir', (45, 32, 45, 37): '"""ref"""', (45, 39, 45, 46): '"""tests"""'}, {}), "(main_dir, 'ref', 'tests')", False, 'from os.path import join\n'), ((48, 10, 48, 39), 'os.path.join', 'join', ({(48, 15, 48, 23): 'main_dir', (48, 25, 48, 30): '"""ref"""', (48, 32, 48, 38): '"""DATA"""'}, {}), "(main_dir, 'ref', 'DATA')", False, 'from os.path import join\n'), ((49, 10, 49, 40), 'os.path.join', 'join', ({(49, 15, 49, 23): 'main_dir', (49, 25, 49, 30): '"""ref"""', (49, 32, 49, 39): '"""utils"""'}, {}), "(main_dir, 'ref', 'utils')", False, 'from os.path import join\n'), ((52, 19, 52, 41), 'os.path.join', 'join', ({(52, 24, 52, 32): 'main_dir', (52, 34, 52, 40): '"""cmts"""'}, {}), "(main_dir, 'cmts')", False, 'from os.path import join\n'), ((53, 26, 53, 60), 'os.path.join', 'join', ({(53, 31, 53, 39): 'main_dir', (53, 41, 53, 47): '"""cmts"""', (53, 49, 53, 59): '"""cmts_raw"""'}, {}), "(main_dir, 'cmts', 'cmts_raw')", False, 'from os.path import join\n'), ((54, 19, 54, 59), 'os.path.join', 'join', ({(54, 24, 54, 32): 'main_dir', (54, 34, 54, 40): '"""cmts"""', (54, 42, 54, 58): '"""cmts_generated"""'}, {}), "(main_dir, 'cmts', 'cmts_generated')", False, 'from os.path import join\n'), ((60, 19, 60, 41), 'os.path.join', 'join', ({(60, 24, 60, 32): 'main_dir', (60, 34, 60, 40): '"""work"""'}, {}), "(main_dir, 'work')", False, 'from os.path import join\n'), ((65, 20, 65, 59), 'os.path.join', 'join', ({(65, 25, 65, 33): 'main_dir', (65, 35, 65, 41): '"""cmts"""', (65, 43, 65, 53): '"""cmts_raw"""', (65, 55, 65, 58): '"""*"""'}, {}), "(main_dir, 'cmts', 'cmts_raw', '*')", False, 'from os.path import join\n'), ((78, 10, 78, 32), 'os.path.join', 'join', ({(78, 15, 78, 23): 'main_dir', (78, 25, 78, 31): '"""DATA"""'}, {}), "(main_dir, 'DATA')", False, 'from os.path import join\n'), ((78, 34, 78, 63), 'os.path.join', 'join', ({(78, 39, 78, 47): 'main_dir', (78, 49, 78, 54): '"""ref"""', (78, 56, 78, 62): '"""DATA"""'}, {}), "(main_dir, 'ref', 'DATA')", False, 'from os.path import join\n'), ((79, 10, 79, 33), 'os.path.join', 'join', ({(79, 15, 79, 23): 'main_dir', (79, 25, 79, 32): '"""utils"""'}, {}), "(main_dir, 'utils')", False, 'from os.path import join\n'), ((79, 35, 79, 65), 'os.path.join', 'join', ({(79, 40, 79, 48): 'main_dir', (79, 50, 79, 55): '"""ref"""', (79, 57, 79, 64): '"""utils"""'}, {}), "(main_dir, 'ref', 'utils')", False, 'from os.path import join\n'), ((109, 19, 109, 52), 'os.path.join', 'join', ({(109, 24, 109, 34): 'output_dir', (109, 36, 109, 51): '"""DATABASES_MPI"""'}, {}), "(output_dir, 'DATABASES_MPI')", False, 'from os.path import join\n'), ((110, 19, 110, 51), 'os.path.join', 'join', ({(110, 24, 110, 34): 'output_dir', (110, 36, 110, 50): '"""OUTPUT_FILES"""'}, {}), "(output_dir, 'OUTPUT_FILES')", False, 'from os.path import join\n'), ((17, 16, 17, 43), 'obspy.read_events', 'obspy.read_events', ({(17, 34, 17, 42): 'cmt_file'}, {}), '(cmt_file)', False, 'import obspy\n'), ((24, 29, 24, 79), 'os.path.join', 'join', ({(24, 34, 24, 52): 'generated_cmts_dir', (24, 54, 24, 69): 'f"""d{depth_per}"""', (24, 71, 24, 78): 'gcmt_id'}, {}), "(generated_cmts_dir, f'd{depth_per}', gcmt_id)", False, 'from os.path import join\n'), ((56, 23, 57, 62), 'os.path.join', 'join', ({(56, 28, 56, 36): 'main_dir', (56, 38, 56, 44): '"""cmts"""', (57, 28, 57, 44): '"""cmts_generated"""', (57, 46, 57, 61): 'f"""d{depth_per}"""'}, {}), "(main_dir, 'cmts', 'cmts_generated', f'd{depth_per}')", False, 'from os.path import join\n'), ((70, 17, 70, 49), 'os.path.join', 'join', ({(70, 22, 70, 30): 'main_dir', (70, 32, 70, 38): '"""work"""', (70, 40, 70, 48): 'cmt_name'}, {}), "(main_dir, 'work', cmt_name)", False, 'from os.path import join\n'), ((26, 31, 26, 58), 'obspy.read_events', 'obspy.read_events', ({(26, 49, 26, 57): 'cmt_file'}, {}), '(cmt_file)', False, 'import obspy\n'), ((74, 24, 74, 45), 'os.path.join', 'join', ({(74, 29, 74, 37): 'main_dir', (74, 39, 74, 44): '"""ref"""'}, {}), "(main_dir, 'ref')", False, 'from os.path import join\n'), ((75, 18, 75, 67), 'os.path.join', 'join', ({(75, 23, 75, 31): 'main_dir', (75, 33, 75, 39): '"""work"""', (75, 41, 75, 49): 'cmt_name', (75, 51, 75, 66): 'f"""d{depth_per}"""'}, {}), "(main_dir, 'work', cmt_name, f'd{depth_per}')", False, 'from os.path import join\n'), ((84, 21, 84, 78), 'os.path.join', 'join', ({(84, 26, 84, 34): 'main_dir', (84, 36, 84, 42): '"""work"""', (84, 44, 84, 52): 'cmt_name', (84, 54, 84, 69): 'f"""d{depth_per}"""', (84, 71, 84, 77): '"""DATA"""'}, {}), "(main_dir, 'work', cmt_name, f'd{depth_per}', 'DATA')", False, 'from os.path import join\n'), ((91, 18, 92, 49), 'os.path.join', 'join', ({(91, 23, 91, 31): 'main_dir', (91, 33, 91, 39): '"""cmts"""', (91, 41, 91, 57): '"""cmts_generated"""', (92, 23, 92, 38): 'f"""d{depth_per}"""', (92, 40, 92, 48): 'cmt_name'}, {}), "(main_dir, 'cmts', 'cmts_generated', f'd{depth_per}', cmt_name)", False, 'from os.path import join\n'), ((92, 51, 92, 123), 'os.path.join', 'join', ({(92, 56, 92, 64): 'main_dir', (92, 66, 92, 72): '"""work"""', (92, 74, 92, 82): 'cmt_name', (92, 84, 92, 99): 'f"""d{depth_per}"""', (92, 101, 92, 107): '"""DATA"""', (92, 109, 92, 122): '"""CMTSOLUTION"""'}, {}), "(main_dir, 'work', cmt_name, f'd{depth_per}', 'DATA', 'CMTSOLUTION')", False, 'from os.path import join\n'), ((93, 18, 93, 59), 'os.path.join', 'join', ({(93, 23, 93, 31): 'main_dir', (93, 33, 93, 38): '"""ref"""', (93, 40, 93, 46): '"""DATA"""', (93, 48, 93, 58): '"""Par_file"""'}, {}), "(main_dir, 'ref', 'DATA', 'Par_file')", False, 'from os.path import join\n'), ((93, 61, 94, 80), 'os.path.join', 'join', ({(94, 16, 94, 24): 'main_dir', (94, 26, 94, 32): '"""work"""', (94, 34, 94, 42): 'cmt_name', (94, 44, 94, 59): 'f"""d{depth_per}"""', (94, 61, 94, 67): '"""DATA"""', (94, 69, 94, 79): '"""Par_file"""'}, {}), "(main_dir, 'work', cmt_name, f'd{depth_per}', 'DATA', 'Par_file')", False, 'from os.path import join\n'), ((95, 18, 95, 59), 'os.path.join', 'join', ({(95, 23, 95, 31): 'main_dir', (95, 33, 95, 38): '"""ref"""', (95, 40, 95, 46): '"""DATA"""', (95, 48, 95, 58): '"""STATIONS"""'}, {}), "(main_dir, 'ref', 'DATA', 'STATIONS')", False, 'from os.path import join\n'), ((95, 61, 96, 80), 'os.path.join', 'join', ({(96, 16, 96, 24): 'main_dir', (96, 26, 96, 32): '"""work"""', (96, 34, 96, 42): 'cmt_name', (96, 44, 96, 59): 'f"""d{depth_per}"""', (96, 61, 96, 67): '"""DATA"""', (96, 69, 96, 79): '"""STATIONS"""'}, {}), "(main_dir, 'work', cmt_name, f'd{depth_per}', 'DATA', 'STATIONS')", False, 'from os.path import join\n'), ((113, 27, 114, 58), 'os.path.join', 'join', ({(113, 32, 113, 42): 'output_dir', (113, 44, 113, 59): '"""DATABASES_MPI"""', (114, 32, 114, 40): 'cmt_name', (114, 42, 114, 57): 'f"""d{depth_per}"""'}, {}), "(output_dir, 'DATABASES_MPI', cmt_name, f'd{depth_per}')", False, 'from os.path import join\n'), ((115, 27, 116, 58), 'os.path.join', 'join', ({(115, 32, 115, 42): 'output_dir', (115, 44, 115, 58): '"""OUTPUT_FILES"""', (116, 32, 116, 40): 'cmt_name', (116, 42, 116, 57): 'f"""d{depth_per}"""'}, {}), "(output_dir, 'OUTPUT_FILES', cmt_name, f'd{depth_per}')", False, 'from os.path import join\n'), ((117, 24, 118, 55), 'os.path.join', 'join', ({(117, 29, 117, 39): 'output_dir', (117, 41, 117, 56): '"""DATABASES_MPI"""', (118, 29, 118, 37): 'cmt_name', (118, 39, 118, 54): 'f"""d{depth_per}"""'}, {}), "(output_dir, 'DATABASES_MPI', cmt_name, f'd{depth_per}')", False, 'from os.path import join\n'), ((118, 57, 118, 123), 'os.path.join', 'join', ({(118, 62, 118, 70): 'main_dir', (118, 72, 118, 78): '"""work"""', (118, 80, 118, 88): 'cmt_name', (118, 90, 118, 105): 'f"""d{depth_per}"""', (118, 107, 118, 122): '"""DATABASES_MPI"""'}, {}), "(main_dir, 'work', cmt_name, f'd{depth_per}', 'DATABASES_MPI')", False, 'from os.path import join\n'), ((119, 24, 120, 55), 'os.path.join', 'join', ({(119, 29, 119, 39): 'output_dir', (119, 41, 119, 55): '"""OUTPUT_FILES"""', (120, 29, 120, 37): 'cmt_name', (120, 39, 120, 54): 'f"""d{depth_per}"""'}, {}), "(output_dir, 'OUTPUT_FILES', cmt_name, f'd{depth_per}')", False, 'from os.path import join\n'), ((120, 57, 120, 122), 'os.path.join', 'join', ({(120, 62, 120, 70): 'main_dir', (120, 72, 120, 78): '"""work"""', (120, 80, 120, 88): 'cmt_name', (120, 90, 120, 105): 'f"""d{depth_per}"""', (120, 107, 120, 121): '"""OUTPUT_FILES"""'}, {}), "(main_dir, 'work', cmt_name, f'd{depth_per}', 'OUTPUT_FILES')", False, 'from os.path import join\n'), ((98, 28, 98, 65), 'os.path.join', 'join', ({(98, 33, 98, 41): 'main_dir', (98, 43, 98, 48): '"""ref"""', (98, 50, 98, 56): '"""DATA"""', (98, 58, 98, 64): 'lnfile'}, {}), "(main_dir, 'ref', 'DATA', lnfile)", False, 'from os.path import join\n'), ((98, 67, 99, 80), 'os.path.join', 'join', ({(99, 20, 99, 28): 'main_dir', (99, 30, 99, 36): '"""work"""', (99, 38, 99, 46): 'cmt_name', (99, 48, 99, 63): 'f"""d{depth_per}"""', (99, 65, 99, 71): '"""DATA"""', (99, 73, 99, 79): 'lnfile'}, {}), "(main_dir, 'work', cmt_name, f'd{depth_per}', 'DATA', lnfile)", False, 'from os.path import join\n'), ((104, 28, 104, 57), 'os.path.join', 'join', ({(104, 33, 104, 41): 'main_dir', (104, 43, 104, 48): '"""ref"""', (104, 50, 104, 56): 'lnfile'}, {}), "(main_dir, 'ref', lnfile)", False, 'from os.path import join\n'), ((104, 59, 105, 72), 'os.path.join', 'join', ({(105, 20, 105, 28): 'main_dir', (105, 30, 105, 36): '"""work"""', (105, 38, 105, 46): 'cmt_name', (105, 48, 105, 63): 'f"""d{depth_per}"""', (105, 65, 105, 71): 'lnfile'}, {}), "(main_dir, 'work', cmt_name, f'd{depth_per}', lnfile)", False, 'from os.path import join\n')]