# -*- coding: utf-8 -*-
"""
OpenMOM (Python)
Version 4.1.0
omm.py : solver
"""

import sys, time
import numba
import sol.input_data, sol.wiregrid, sol.solve, sol.calcchar
import sol.monitor, sol.output_log, sol.plot3d_geom, sol.save_bin

def main(argv):
    # ロゴ
    version = 'OpenMOM (Python) Version 4.1.0'

    # CPU/GPU
    GPU = 1  # 0=CPU, 1=GPU

    # CPUスレッド数
    thread = 8

    # 入力データファイル名(omm.pyからの相対パス)
    omm_in = 'python.omm'
    #omm_in = '../_tmp.omm'
    #omm_in = '../data/sample/1st_sample.omm'
    #omm_in = '../data/benchmark/n5000.omm'

    # 単精度/倍精度
    c_dtype = 'c8'  # 単精度='c8', 倍精度='c16'

    # 終了時のプロンプト
    prompt = 0

    # 引数処理(引数が優先される)
    if len(argv) > 1:
        GPU, thread, c_dtype, prompt, omm_in = _args(argv, version)

    # 出力ファイル名(固定)
    fn_log = 'omm.log'
    fn_out = 'omm.npz'

    # Numbaスレッド数設定
    numba.set_num_threads(thread)

    # ロゴ
    logo = '<<< %s >>>\nCPU-thread=%d, GPU=%s, precision=%s' % \
      (version, thread, ('on' if GPU else 'off'), ('single' if c_dtype == 'c8' else 'double'))

    # cpu time
    cpu = [0] * 5
    cpu[0] = time.time()

    # logファイルを開く
    fp_log = open(fn_log, 'wt', encoding='utf-8')

    # 経過確認 (1)
    sol.monitor.monitor1(fp_log, logo)

    # [1] データ入力
    Ngeom, G_gtype, G_cosys, G_pos, G_div, G_ifeed, G_feed, G_iload, G_load, G_iradius, G_radius, G_offset, \
    Nfreq, Freq, Title, Iplanewave, Planewave, Iground, Iradiusall, Radiusall, Z0, Plot3dgeom \
        = sol.input_data.read(omm_in)

    # [2] ワイヤグリッドモデルを作成する
    Ne, Nfeed, Nload, E_posc, E_posm, E_posp, E_lng, E_tan, E_rad, E_ifeed, E_feed, E_iload, E_load \
    = sol.wiregrid.makedata(
        Iradiusall, Radiusall, Ngeom, G_gtype, G_cosys, G_pos, G_div, G_ifeed, G_feed, G_iload, G_load, G_iradius, G_radius, G_offset, Iground)
    sol.output_log.element(Ne, E_posc, E_lng, E_tan, E_rad)

    # 3D図形表示で確認する
    if Plot3dgeom == 1:
        sol.plot3d_geom.plot(Title, Ngeom, Ne, E_posm, E_posp, E_ifeed, E_iload)

    # 経過確認 (2)
    sol.monitor.monitor2(fp_log, c_dtype, Title, Ngeom, Ne, Nfreq, Nfeed, Nload, Iplanewave, Iground)
    
    cpu[1] = time.time()

    # [3] 電流分布を計算する(計算の主要部)
    Xc, tcpu = sol.solve.solve(GPU, c_dtype, Ne, Nfreq, Freq, Nfeed, Iplanewave, Planewave,
        E_posc, E_posm, E_posp, E_lng, E_tan, E_rad, E_ifeed, E_feed, E_iload, E_load, Iground)
    sol.output_log.current(Ne, Nfreq, E_ifeed, E_iload, Xc)

    cpu[2] = cpu[1] + tcpu[0]
    cpu[3] = cpu[2] + tcpu[1]

    # [4] 入力インピーダンスと散乱断面積を出力する
    Zin, Rcs = sol.calcchar.calc(
        Ne, Nfreq, Freq, Nfeed, E_posc, E_lng, E_tan, E_ifeed, E_feed, Planewave, Iground, Z0, Xc)
    
    # 経過確認 (3)
    sol.monitor.monitor3(fp_log, Nfreq, Freq, Nfeed, Z0, Zin, Rcs)

    # 経過確認 (4)
    sol.monitor.monitor4(fp_log)

    # [5] 計算結果をファイルに保存する
    sol.save_bin.save(fn_out,
        Ne, Nfreq, Freq, Nfeed, Iplanewave, Planewave,
        E_posc, E_posm, E_posp, E_lng, E_tan, E_ifeed, E_feed, E_iload, E_load,
        Iground, Title, Z0, Xc, Zin)

    # 経過確認 (5)
    cpu[4] = time.time()
    sol.monitor.monitor5(fp_log, cpu)

    # logファイルを閉じる
    fp_log.close()

    # メモリー解放
    Xc = None

    # prompt
    if prompt:
        input()

# (private) 引数処理
def _args(argv, version):

    usage = 'Usage : python omm.py [-cpu|-gpu] [-n <thread>] [-single|-double] <datafile>'
    GPU = 0
    thread = 1
    c_dtype = 'c8'
    prompt = 0
    omm_in = ''

    i = 1
    while i < len(argv):
        arg = argv[i].lower()
        if   arg == '-gpu':
            GPU = 1
            i += 1
        elif arg == '-cpu':
            GPU = 0
            i += 1
        elif arg == '-n':
            thread = int(argv[i + 1])
            i += 2
        elif arg == '-single':
            c_dtype = 'c8'
            i += 1
        elif arg == '-double':
            c_dtype = 'c16'
            i += 1
        elif arg == '-prompt':
            prompt = 1
            i += 1
        elif arg == '--help':
            print(usage)
            sys.exit()
        elif arg == '--version':
            print(version)
            sys.exit()
        else:
            omm_in = argv[i]
            i += 1

    return GPU, thread, c_dtype, prompt, omm_in

# entry point
if __name__ == "__main__":
    main(sys.argv)
