# -*- coding: utf-8 -*-

"""Definition of unit systems in LAMMPS."""
from seamm_util import ureg, Q_, units_class  # noqa: F401

lammps_unit_system = 'real'

unit_systems = {
    'real': {
        'mass': 'grams/mole',
        'distance': 'angstroms',
        'time': 'femtoseconds',
        'energy': 'kcal/mole',
        'velocity': 'angstroms/femtosecond',
        'force': 'kcal/mole/angstrom',
        'torque': 'kcal/mole',
        'temperature': 'kelvin',
        'pressure': 'atmospheres',
        'dynamic viscosity': 'poise',
        'charge': 'elementary_charge',
        'dipole': 'elementary_charge*angstroms',
        'electric field': 'volts/angstrom',
        'density': 'gram/cm^3',
    },
    'metal': {
        'mass': 'grams/mole',
        'distance': 'angstroms',
        'time': 'picoseconds',
        'energy': 'eV',
        'velocity': 'angstroms/picosecond',
        'force': 'eV/angstrom',
        'torque': 'eV',
        'temperature': 'kelvin',
        'pressure': 'bars',
        'dynamic viscosity': 'poise',
        'charge': 'elementary_charge',
        'dipole': 'elementary_charge*angstroms',
        'electric field': 'volts/angstrom',
        'density': 'gram/cm^3',
    },
    'si': {
        'mass': 'kilograms',
        'distance': 'meters',
        'time': 'seconds',
        'energy': 'joules',
        'velocity': 'meters/second',
        'force': 'newtons',
        'torque': 'newton*meters',
        'temperature': 'kelvin',
        'pressure': 'pascals',
        'dynamic viscosity': 'pascal*second',
        'charge': 'coulombs',
        'dipole': 'coulombs*meters',
        'electric field': 'volts/meter',
        'density': 'kilograms/meter^3',
    },
    'cgs': {
        'mass': 'grams',
        'distance': 'centimeters',
        'time': 'seconds',
        'energy': 'ergs',
        'velocity': 'centimeters/second',
        'force': 'dynes',
        'torque': 'dyne*centimeters',
        'temperature': 'kelvin',
        'pressure': 'dyne/cm^2',
        'dynamic viscosity': 'poise',
        'charge': 'statcoulombs',
        'dipole': 'statcoulombs*cm',
        'electric field': 'statvolt/cm',
        'density': 'grams/cm^3',
    },
    'electron': {
        'mass': 'amu',
        'distance': 'bohr',
        'time': 'femtoseconds',
        'energy': 'hartrees',
        'velocity': 'bohr/a_u_time',
        'force': 'hartrees/bohr',
        'temperature': 'kelvin',
        'pressure': 'pascals',
        'charge': 'elementary_charge',
        'dipole moment': 'debye',
        'electric field': 'volts/cm',
    },
    'micro': {
        'mass': 'picograms',
        'distance': 'micrometers',
        'time': 'microseconds',
        'energy': 'picogram*micrometer^2/microsecond^2',
        'velocity': 'micrometers/microsecond',
        'force': 'picogram*micrometer/microsecond^2',
        'torque': 'picogram*micrometer^2/microsecond^2',
        'temperature': 'kelvin',
        'pressure': 'picogram/(micrometer*microsecond^2)',
        'dynamic viscosity': 'picogram/(micrometer*microsecond)',
        'charge': 'picocoulombs',
        'dipole': 'picocoulomb*micrometer',
        'electric field': 'volt/micrometer',
        'density': 'picograms/micrometer^3',
    },
    'nano': {
        'mass': 'attograms',
        'distance': 'nanometers',
        'time': 'nanoseconds',
        'energy': 'attogram*nanometer^2/nanosecond^2',
        'velocity': 'nanometers/nanosecond',
        'force': 'attogram*nanometer/nanosecond^2',
        'torque': 'attogram*nanometer^2/nanosecond^2',
        'temperature': 'kelvin',
        'pressure': 'attogram/(nanometer*nanosecond^2)',
        'dynamic viscosity': 'attogram/(nanometer*nanosecond)',
        'charge': 'elementary_charge',
        'dipole': 'elementary_charge*nanometer',
        'electric field': 'volt/nanometer',
        'density': 'attograms/nanometer^3',
    }
}

quantity_to_dimensionality = {
    "mass": [
        "[mass] / [substance]",
        "[mass]"
    ],
    "distance": [
        "[length]"
    ],
    "time": [
        "[time]"
    ],
    "energy": [
        "[length] ** 2 * [mass] / [substance] / [time] ** 2",
        "[length] ** 2 * [mass] / [time] ** 2"
    ],
    "velocity": [
        "[length] / [time]"
    ],
    "force": [
        "[length] * [mass] / [substance] / [time] ** 2",
        "[length] * [mass] / [time] ** 2"
    ],
    "torque": [
        "[length] ** 2 * [mass] / [substance] / [time] ** 2",
        "[length] ** 2 * [mass] / [time] ** 2"
    ],
    "temperature": [
        "[temperature]"
    ],
    "pressure": [
        "[mass] / [length] / [time] ** 2"
    ],
    "dynamic viscosity": [
        "[mass] / [length] / [time]"
    ],
    "charge": [
        "[current] * [time]",
        "[length] ** 1.5 * [mass] ** 0.5 / [time]"
    ],
    "dipole": [
        "[current] * [length] * [time]",
        "[length] ** 2.5 * [mass] ** 0.5 / [time]"
    ],
    "electric field": [
        "[length] * [mass] / [current] / [time] ** 3",
        "[mass] ** 0.5 / [length] ** 0.5 / [time]"
    ],
    "density": [
        "[mass] / [length] ** 3"
    ],
    "dipole moment": [
        "[current] * [length] * [time]"
    ]
}

dimensionality_to_quantity = {
    "[mass] / [substance]": [
        "mass"
    ],
    "[mass]": [
        "mass"
    ],
    "[length]": [
        "distance"
    ],
    "[time]": [
        "time"
    ],
    "[length] ** 2 * [mass] / [substance] / [time] ** 2": [
        "energy",
        "torque"
    ],
    "[length] ** 2 * [mass] / [time] ** 2": [
        "energy",
        "torque"
    ],
    "[length] / [time]": [
        "velocity"
    ],
    "[length] * [mass] / [substance] / [time] ** 2": [
        "force"
    ],
    "[length] * [mass] / [time] ** 2": [
        "force"
    ],
    "[temperature]": [
        "temperature"
    ],
    "[mass] / [length] / [time] ** 2": [
        "pressure"
    ],
    "[mass] / [length] / [time]": [
        "dynamic viscosity"
    ],
    "[current] * [time]": [
        "charge"
    ],
    "[length] ** 1.5 * [mass] ** 0.5 / [time]": [
        "charge"
    ],
    "[current] * [length] * [time]": [
        "dipole",
        "dipole moment"
    ],
    "[length] ** 2.5 * [mass] ** 0.5 / [time]": [
        "dipole"
    ],
    "[length] * [mass] / [current] / [time] ** 3": [
        "electric field"
    ],
    "[mass] ** 0.5 / [length] ** 0.5 / [time]": [
        "electric field"
    ],
    "[mass] / [length] ** 3": [
        "density"
    ]
}


def set_lammps_unit_system(value):
    """Set the default unit system for LAMMPS."""
    global lammps_unit_system
    if value not in unit_systems:
        raise ValueError(f"unit system '{value}' not recognized!")
    lammps_unit_system = value


def get_lammps_unit_system(value):
    """Get the default unit system for LAMMPS."""
    global lammps_unit_system
    return lammps_unit_system


def to_lammps_units(value, units=None, quantity=None, unit_system=None):
    """Convert the value with units to the current LAMMPS units."""
    if not isinstance(value, units_class):
        if units is not None:
            value = Q_(value, units)
        else:
            value = Q_(value)

    if quantity is None:
        quantities = dimensionality_to_quantity[str(value.dimensionality)]
        quantity = quantities[0]

    if unit_system is None:
        unit_system = lammps_unit_system

    unit_expr = unit_systems[unit_system][quantity]
    converted = value.to(unit_expr)

    return converted.magnitude


def from_lammps_units(value, units, quantity=None, unit_system=None):
    """Convert the scalar value from LAMMPS units to requested units"""
    if quantity is None:
        dimensions = Q_(1.0, units).dimensionality
        quantities = dimensionality_to_quantity[str(dimensions)]
        quantity = quantities[0]

    if unit_system is None:
        unit_system = lammps_unit_system

    unit_expr = unit_systems[unit_system][quantity]
    converted = Q_(value, unit_expr).to(units)

    return converted


if __name__ == '__main__':  # pragma: no cover
    import pint
    import json

    for unit in ureg:
        try:
            q = ureg(unit)
        except pint.errors.UndefinedUnitError:
            q = 'unknown'
        print(f'{unit} --> {q}')
    print('')

    dimensionality = {}
    for system, units in unit_systems.items():
        print(f'{system} system of units')
        for quantity, unit_expr in units.items():
            print(f'\t{quantity} uses {unit_expr}')
            q = ureg(unit_expr)
            print(f'\t\t{q}')
            dimensions = str(q.dimensionality)
            if quantity not in dimensionality:
                dimensionality[quantity] = [dimensions]
            elif dimensions not in dimensionality[quantity]:
                dimensionality[quantity].append(dimensions)
                print(f'{system}: {quantity} <-- {unit_expr}')

    print()
    print(json.dumps(dimensionality, indent=4))

    inverse = {}
    for quantity, unit_exprs in dimensionality.items():
        for unit_expr in unit_exprs:
            if unit_expr not in inverse:
                inverse[unit_expr] = [quantity]
            elif quantity != inverse[unit_expr]:
                inverse[unit_expr].append(quantity)
                print(f'quantity {quantity} != {inverse[unit_expr]} for {unit_expr}')  # noqa: E501

    print()
    print(json.dumps(inverse, indent=4))

    print()
    value = '2.0 kJ/mol'
    converted = to_lammps_units(value)
    print(f'{value} in real units is {converted:.4}')
    print(f'{converted:.4} --> {from_lammps_units(converted, "kJ/mol"):~.4}')

    lammps_unit_system = 'metal'
    converted = to_lammps_units(value)
    print(f'{value} in metal units is {converted:.4}')
    print(f'{converted:.4} --> {from_lammps_units(converted, "kJ/mol"):~.4}')
