import argparse
import multiprocessing
import asyncio
import time

from ..calculator import LammpsLocalClient


script = """
log  lammps.log
units  metal
dimension  3
boundary  p p p
atom_style  full
read_data  initial.data
kspace_style  pppm 0.000010
pair_style  buck/coul/long 10
pair_coeff 1 1 1309362.2766468062 0.104 0.0
pair_coeff 1 2 9892.357 0.20199 0.0
pair_coeff 2 2 2145.7345 0.3 30.2222
set type 2 charge -1.400000
set type 1 charge 1.400000
dump  1 all custom 1 mol.lammpstrj id type x y z fx fy fz
dump_modify  1 sort id
thermo_style  custom step etotal pxx pyy pzz pxy pxz pyz
run  0
"""

files = {
    'initial.data': """pymatgen autogenerated data file

8 atoms

2 atom types

0 4.1990858 xlo xhi
0 4.1990858 ylo yhi
0 4.1990858 zlo zhi
2.571198492157551e-16 2.571198492157551e-16 2.5711984921575517e-16 xy xz yz

Masses

1 24.305
2 15.9994

Atoms

1 1 1 1.4 0.0 0.0 0.0
2 1 1 1.4 3.3763224914577695e-16 2.0995429 2.0995429
3 1 1 1.4 2.0995429 0.0 2.0995429
4 1 1 1.4 2.0995429000000003 2.0995429 2.571198492157551e-16
5 1 2 -1.4 2.0995429000000003 2.0995429 2.0995429000000003
6 1 2 -1.4 2.0995429 0.0 1.2855992460787756e-16
7 1 2 -1.4 3.3763224914577695e-16 2.0995429 1.2855992460787756e-16
8 1 2 -1.4 0.0 0.0 2.0995429
"""
}


def add_subcommand_benchmark(subparsers):
    parser = subparsers.add_parser('benchmark', help='benchmark lammps package')
    parser.set_defaults(func=handle_subcommand_benchmark)
    parser.add_argument('test', choices={'local_client',}, help='benchmark to run')
    parser.add_argument('--max-workers', type=int, default=multiprocessing.cpu_count())
    parser.add_argument('--num-tasks', type=int, default=1000)
    parser.add_argument('--command', default='lammps_serial')


def handle_subcommand_benchmark(args):
    max_workers = min([multiprocessing.cpu_count(), args.max_workers])

    async def run_lammps_job(client, num_workers, num_tasks):
        jobs = []
        start_time = time.perf_counter()
        for _ in range(num_tasks):
            jobs.append(await client.submit(script, files, properties={'stress', 'forces', 'energy'}))
        results = await asyncio.gather(*jobs)
        end_time = time.perf_counter()
        print('LammpsLocalClient num_workers', num_workers, 'time', num_tasks, end_time - start_time, 'tasks/sec', num_tasks / (end_time - start_time), 'tasks/(sec worker)', num_tasks / (end_time - start_time) / num_workers)
        return results

    if args.test == 'local_client':
        loop = asyncio.get_event_loop()

        for num_workers in range(1, max_workers+1):
            try:
                client = LammpsLocalClient(command=args.command, num_workers=num_workers)
                loop.run_until_complete(client.create())
                loop.run_until_complete(run_lammps_job(client, num_workers, args.num_tasks))
            finally:
                client.shutdown()

import argparse
import multiprocessing
import asyncio
import time

from ..calculator import LammpsLocalClient


script = """
log  lammps.log
units  metal
dimension  3
boundary  p p p
atom_style  full
read_data  initial.data
kspace_style  pppm 0.000010
pair_style  buck/coul/long 10
pair_coeff 1 1 1309362.2766468062 0.104 0.0
pair_coeff 1 2 9892.357 0.20199 0.0
pair_coeff 2 2 2145.7345 0.3 30.2222
set type 2 charge -1.400000
set type 1 charge 1.400000
dump  1 all custom 1 mol.lammpstrj id type x y z fx fy fz
dump_modify  1 sort id
thermo_style  custom step etotal pxx pyy pzz pxy pxz pyz
run  0
"""

files = {
    'initial.data': """pymatgen autogenerated data file

8 atoms

2 atom types

0 4.1990858 xlo xhi
0 4.1990858 ylo yhi
0 4.1990858 zlo zhi
2.571198492157551e-16 2.571198492157551e-16 2.5711984921575517e-16 xy xz yz

Masses

1 24.305
2 15.9994

Atoms

1 1 1 1.4 0.0 0.0 0.0
2 1 1 1.4 3.3763224914577695e-16 2.0995429 2.0995429
3 1 1 1.4 2.0995429 0.0 2.0995429
4 1 1 1.4 2.0995429000000003 2.0995429 2.571198492157551e-16
5 1 2 -1.4 2.0995429000000003 2.0995429 2.0995429000000003
6 1 2 -1.4 2.0995429 0.0 1.2855992460787756e-16
7 1 2 -1.4 3.3763224914577695e-16 2.0995429 1.2855992460787756e-16
8 1 2 -1.4 0.0 0.0 2.0995429
"""
}


def add_subcommand_benchmark(subparsers):
    parser = subparsers.add_parser('benchmark', help='benchmark lammps package')
    parser.set_defaults(func=handle_subcommand_benchmark)
    parser.add_argument('test', choices={'local_client',}, help='benchmark to run')
    parser.add_argument('--max-workers', type=int, default=multiprocessing.cpu_count())
    parser.add_argument('--num-tasks', type=int, default=1000)
    parser.add_argument('--command', default='lammps_serial')


def handle_subcommand_benchmark(args):
    max_workers = min([multiprocessing.cpu_count(), args.max_workers])

    async def run_lammps_job(client, num_workers, num_tasks):
        jobs = []
        start_time = time.perf_counter()
        for _ in range(num_tasks):
            jobs.append(await client.submit(script, files, properties={'stress', 'forces', 'energy'}))
        results = await asyncio.gather(*jobs)
        end_time = time.perf_counter()
        print('LammpsLocalClient num_workers', num_workers, 'time', num_tasks, end_time - start_time, 'tasks/sec', num_tasks / (end_time - start_time), 'tasks/(sec worker)', num_tasks / (end_time - start_time) / num_workers)
        return results

    if args.test == 'local_client':
        loop = asyncio.get_event_loop()

        for num_workers in range(1, max_workers+1):
            try:
                client = LammpsLocalClient(command=args.command, num_workers=num_workers)
                loop.run_until_complete(client.create())
                loop.run_until_complete(run_lammps_job(client, num_workers, args.num_tasks))
            finally:
                client.shutdown()

log  lammps.log units  metal dimension  3 boundary  p p p atom_style  full read_data  initial.data kspace_style  pppm 0.000010 pair_style  buck/coul/long 10 pair_coeff 1 1 1309362.2766468062 0.104 0.0 pair_coeff 1 2 9892.357 0.20199 0.0 pair_coeff 2 2 2145.7345 0.3 30.2222 set type 2 charge -1.400000 set type 1 charge 1.400000 dump  1 all custom 1 mol.lammpstrj id type x y z fx fy fz dump_modify  1 sort id thermo_style  custom step etotal pxx pyy pzz pxy pxz pyz run  0 pymatgen autogenerated data file  8 atoms  2 atom types  0 4.1990858 xlo xhi 0 4.1990858 ylo yhi 0 4.1990858 zlo zhi 2.571198492157551e-16 2.571198492157551e-16 2.5711984921575517e-16 xy xz yz  Masses  1 24.305 2 15.9994  Atoms  1 1 1 1.4 0.0 0.0 0.0 2 1 1 1.4 3.3763224914577695e-16 2.0995429 2.0995429 3 1 1 1.4 2.0995429 0.0 2.0995429 4 1 1 1.4 2.0995429000000003 2.0995429 2.571198492157551e-16 5 1 2 -1.4 2.0995429000000003 2.0995429 2.0995429000000003 6 1 2 -1.4 2.0995429 0.0 1.2855992460787756e-16 7 1 2 -1.4 3.3763224914577695e-16 2.0995429 1.2855992460787756e-16 8 1 2 -1.4 0.0 0.0 2.0995429