from impedance.models.circuits import BaseCircuit, CustomCircuit, Randles
import json
import matplotlib.pyplot as plt
import numpy as np
import os
import pytest

# get example data
data = np.genfromtxt(os.path.join("./data/", "exampleData.csv"), delimiter=",")

f = data[:, 0]
Z = data[:, 1] + 1j * data[:, 2]


def test_BaseCircuit():
    initial_guess = [0.01, 0.02, 50]
    base_circuit = BaseCircuit(initial_guess)

    # __init__()
    # check initial_guess is loaded in correctly
    assert base_circuit.initial_guess == initial_guess

    # improper input_guess types raise an TypeError
    with pytest.raises(TypeError):
        r = BaseCircuit(initial_guess=["hi", 0.1])

    # __eq__()
    # incorrect comparisons raise a TypeError
    with pytest.raises(TypeError):
        r = BaseCircuit(initial_guess=[0.01, 0.005, 0.1, 0.0001, 200])
        r == 8

    # fit()
    # improper data types in fitting raise a TypeError
    with pytest.raises(TypeError):
        r = BaseCircuit(initial_guess=[0.01, 0.005, 0.1, 0.0001, 200])
        r.fit([42, 4.2], [])  # frequencies not ndarray

    with pytest.raises(TypeError):
        r = BaseCircuit(initial_guess=[0.01, 0.005, 0.1, 0.0001, 200])
        r.fit(np.array([42 + 42j]), [])  # frequencies not numeric type

    with pytest.raises(TypeError):
        r = BaseCircuit(initial_guess=[0.01, 0.005, 0.1, 0.0001, 200])
        r.fit(np.array([42]), [42 + 42j])  # Z not ndarray

    with pytest.raises(TypeError):
        r = BaseCircuit(initial_guess=[0.01, 0.005, 0.1, 0.0001, 200])
        r.fit(np.array([42]), np.array([0.5, 0.2]))  # Z not complex

    with pytest.raises(TypeError):
        r = BaseCircuit(initial_guess=[0.01, 0.005, 0.1, 0.0001, 200])
        r.fit(np.array([42, 4.2]), np.array([42 + 42j]))  # mismatched lengths

    # predict()
    # improper data types in fitting raise a TypeError
    with pytest.raises(TypeError):
        r = BaseCircuit(initial_guess=[0.01, 0.005, 0.1, 0.0001, 200])
        r.predict([42, 4.2])  # frequencies not ndarray

    with pytest.raises(TypeError):
        r = BaseCircuit(initial_guess=[0.01, 0.005, 0.1, 0.0001, 200])
        r.predict(np.array([42 + 42j]))  # frequencies not numeric type

    # plot()
    # kind = {'nyquist', 'bode'} should return a plt.Axes() object
    _, ax = plt.subplots()
    assert isinstance(base_circuit.plot(ax, None, Z, kind="nyquist"), type(ax))
    assert isinstance(base_circuit.plot(None, f, Z, kind="nyquist"), type(ax))
    _, axes = plt.subplots(nrows=2)
    assert isinstance(base_circuit.plot(axes, f, Z, kind="bode")[0], type(ax))
    assert isinstance(base_circuit.plot(None, f, Z, kind="bode")[0], type(ax))

    # incorrect kind raises a ValueError
    with pytest.raises(ValueError):
        base_circuit.plot(None, f, Z, kind="SomethingElse")


def test_Randles():
    randles = Randles(initial_guess=[0.01, 0.005, 0.1, 0.01, 200])
    randlesCPE = Randles(initial_guess=[0.01, 0.05, 0.1, 0.9, 0.01, 200], CPE=True)
    with pytest.raises(ValueError):
        randlesCPE = Randles([0.01, 200])  # incorrect initial guess length
    randles.fit(f[np.imag(Z) < 0], Z[np.imag(Z) < 0])
    randlesCPE.fit(f[np.imag(Z) < 0], Z[np.imag(Z) < 0])

    # compare with known fit parameters
    np.testing.assert_almost_equal(
        randles.parameters_,
        np.array(
            [
                1.86146620e-02,
                1.15477171e-02,
                1.33331949e00,
                6.31473571e-02,
                2.22407275e02,
            ]
        ),
        decimal=2,
    )

    # compare with known impedance predictions
    assert np.isclose(
        randles.predict(np.array([10.0])), np.complex(0.02495749, -0.00614842)
    )

    # check altair plotting with a fit circuit
    chart = randles.plot(f_data=f, Z_data=Z)
    datasets = json.loads(chart.to_json())["datasets"]
    for dataset in datasets.keys():
        assert len(datasets[dataset]) == len(Z)

    # plot() with fitted model
    # check defaults work if no frequency data is passed
    chart = randles.plot(Z_data=Z)

    # bode plots
    randles.plot(f_data=f, Z_data=Z, kind="bode")
    randles.plot(kind="bode")
    with pytest.raises(ValueError):
        randles.plot(Z_data=Z, kind="bode")  # missing f_data

    # nyquist plots
    randles.plot(f_data=f, Z_data=Z, kind="nyquist")
    randles.plot(Z_data=Z, kind="nyquist")

    # check equality comparisons work
    randles1 = Randles(initial_guess=[0.01, 0.005, 0.1, 0.0001, 200])
    randles2 = Randles(initial_guess=[0.01, 0.005, 0.1, 0.0001, 200])
    assert randles1 == randles2

    randles1.fit(f[np.imag(Z) < 0], Z[np.imag(Z) < 0])
    assert randles1 != randles2

    randles2.fit(f[np.imag(Z) < 0], Z[np.imag(Z) < 0])
    assert randles1 == randles2

    randles2.fit(f, Z)
    assert randles1 != randles2


def test_CustomCircuit():
    initial_guess = [0.01, 0.005, 0.1, 0.005, 0.1, 0.001, 200]
    custom_string = "R0-p(R1,C1)-p(R2,C2)-Wo1"
    custom_circuit = CustomCircuit(initial_guess=initial_guess, circuit=custom_string)

    # check get_param_names()
    full_names, all_units = custom_circuit.get_param_names()
    assert full_names == ["R0", "R1", "C1", "R2", "C2", "Wo1_0", "Wo1_1"]
    assert all_units == ["Ohm", "Ohm", "F", "Ohm", "F", "Ohm", "sec"]

    # check _is_fit()
    assert not custom_circuit._is_fit()

    # check predictions from initial_guesses
    high_f = np.array([1e9])
    assert np.isclose(np.real(custom_circuit.predict(high_f)[0]), initial_guess[0])

    # __str()__
    initial_guess = [0.01, 0.005, 0.1]
    custom_string = "R0-p(R1,C1)"
    custom_circuit = CustomCircuit(initial_guess=initial_guess, circuit=custom_string)

    assert (
        str(custom_circuit)
        == "\nCircuit string: R0-p(R1,C1)\n"
        + "Fit: False\n"
        + "\nInitial guesses:\n"
        + "     R0 = 1.00e-02 [Ohm]\n"
        + "     R1 = 5.00e-03 [Ohm]\n"
        + "     C1 = 1.00e-01 [F]\n"
    )

    custom_circuit.fit(f, Z)
    assert custom_circuit._is_fit()
    custom_circuit.plot(f_data=f, Z_data=Z)

    # constants and _ in circuit and no name
    circuit = "R_0-p(R_1,C_1)-Wo_1"
    constants = {"R_0": 0.02, "Wo_1_1": 200}
    initial_guess = [0.005, 0.1, 0.001]
    custom_circuit = CustomCircuit(
        initial_guess=initial_guess, constants=constants, circuit=circuit, name="Test"
    )

    assert (
        str(custom_circuit)
        == "\nName: Test\n"
        + "Circuit string: R_0-p(R_1,C_1)-Wo_1\n"
        + "Fit: False\n"
        + "\nConstants:\n"
        + "    R_0 = 2.00e-02 [Ohm]\n"
        + "  Wo_1_1 = 2.00e+02 [sec]\n"
        + "\nInitial guesses:\n"
        + "    R_1 = 5.00e-03 [Ohm]\n"
        + "    C_1 = 1.00e-01 [F]\n"
        + "  Wo_1_0 = 1.00e-03 [Ohm]\n"
    )

    # incorrect number of initial guesses
    with pytest.raises(ValueError):
        initial_guess = [0.01, 0.005, 0.1, 0.005, 0.1, 0.001, 200]
        custom_string = "R0-p(R1,CPE1)-p(R1,C1)-Wo1"
        custom_circuit = CustomCircuit(
            initial_guess=initial_guess, circuit=custom_string
        )

    # no initial guesses supplied before fitting
    with pytest.raises(ValueError):
        custom_circuit = CustomCircuit()
        custom_circuit.fit(f, Z)

    # incorrect circuit element in circuit
    with pytest.raises(ValueError):
        custom_circuit = CustomCircuit("R0-NotAnElement", initial_guess=[1, 2])

from impedance.models.circuits import BaseCircuit, CustomCircuit, Randles
import json
import matplotlib.pyplot as plt
import numpy as np
import os
import pytest

# get example data
data = np.genfromtxt(os.path.join("./data/", "exampleData.csv"), delimiter=",")

f = data[:, 0]
Z = data[:, 1] + 1j * data[:, 2]


def test_BaseCircuit():
    initial_guess = [0.01, 0.02, 50]
    base_circuit = BaseCircuit(initial_guess)

    # __init__()
    # check initial_guess is loaded in correctly
    assert base_circuit.initial_guess == initial_guess

    # improper input_guess types raise an TypeError
    with pytest.raises(TypeError):
        r = BaseCircuit(initial_guess=["hi", 0.1])

    # __eq__()
    # incorrect comparisons raise a TypeError
    with pytest.raises(TypeError):
        r = BaseCircuit(initial_guess=[0.01, 0.005, 0.1, 0.0001, 200])
        r == 8

    # fit()
    # improper data types in fitting raise a TypeError
    with pytest.raises(TypeError):
        r = BaseCircuit(initial_guess=[0.01, 0.005, 0.1, 0.0001, 200])
        r.fit([42, 4.2], [])  # frequencies not ndarray

    with pytest.raises(TypeError):
        r = BaseCircuit(initial_guess=[0.01, 0.005, 0.1, 0.0001, 200])
        r.fit(np.array([42 + 42j]), [])  # frequencies not numeric type

    with pytest.raises(TypeError):
        r = BaseCircuit(initial_guess=[0.01, 0.005, 0.1, 0.0001, 200])
        r.fit(np.array([42]), [42 + 42j])  # Z not ndarray

    with pytest.raises(TypeError):
        r = BaseCircuit(initial_guess=[0.01, 0.005, 0.1, 0.0001, 200])
        r.fit(np.array([42]), np.array([0.5, 0.2]))  # Z not complex

    with pytest.raises(TypeError):
        r = BaseCircuit(initial_guess=[0.01, 0.005, 0.1, 0.0001, 200])
        r.fit(np.array([42, 4.2]), np.array([42 + 42j]))  # mismatched lengths

    # predict()
    # improper data types in fitting raise a TypeError
    with pytest.raises(TypeError):
        r = BaseCircuit(initial_guess=[0.01, 0.005, 0.1, 0.0001, 200])
        r.predict([42, 4.2])  # frequencies not ndarray

    with pytest.raises(TypeError):
        r = BaseCircuit(initial_guess=[0.01, 0.005, 0.1, 0.0001, 200])
        r.predict(np.array([42 + 42j]))  # frequencies not numeric type

    # plot()
    # kind = {'nyquist', 'bode'} should return a plt.Axes() object
    _, ax = plt.subplots()
    assert isinstance(base_circuit.plot(ax, None, Z, kind="nyquist"), type(ax))
    assert isinstance(base_circuit.plot(None, f, Z, kind="nyquist"), type(ax))
    _, axes = plt.subplots(nrows=2)
    assert isinstance(base_circuit.plot(axes, f, Z, kind="bode")[0], type(ax))
    assert isinstance(base_circuit.plot(None, f, Z, kind="bode")[0], type(ax))

    # incorrect kind raises a ValueError
    with pytest.raises(ValueError):
        base_circuit.plot(None, f, Z, kind="SomethingElse")


def test_Randles():
    randles = Randles(initial_guess=[0.01, 0.005, 0.1, 0.01, 200])
    randlesCPE = Randles(initial_guess=[0.01, 0.05, 0.1, 0.9, 0.01, 200], CPE=True)
    with pytest.raises(ValueError):
        randlesCPE = Randles([0.01, 200])  # incorrect initial guess length
    randles.fit(f[np.imag(Z) < 0], Z[np.imag(Z) < 0])
    randlesCPE.fit(f[np.imag(Z) < 0], Z[np.imag(Z) < 0])

    # compare with known fit parameters
    np.testing.assert_almost_equal(
        randles.parameters_,
        np.array(
            [
                1.86146620e-02,
                1.15477171e-02,
                1.33331949e00,
                6.31473571e-02,
                2.22407275e02,
            ]
        ),
        decimal=2,
    )

    # compare with known impedance predictions
    assert np.isclose(
        randles.predict(np.array([10.0])), np.complex(0.02495749, -0.00614842)
    )

    # check altair plotting with a fit circuit
    chart = randles.plot(f_data=f, Z_data=Z)
    datasets = json.loads(chart.to_json())["datasets"]
    for dataset in datasets.keys():
        assert len(datasets[dataset]) == len(Z)

    # plot() with fitted model
    # check defaults work if no frequency data is passed
    chart = randles.plot(Z_data=Z)

    # bode plots
    randles.plot(f_data=f, Z_data=Z, kind="bode")
    randles.plot(kind="bode")
    with pytest.raises(ValueError):
        randles.plot(Z_data=Z, kind="bode")  # missing f_data

    # nyquist plots
    randles.plot(f_data=f, Z_data=Z, kind="nyquist")
    randles.plot(Z_data=Z, kind="nyquist")

    # check equality comparisons work
    randles1 = Randles(initial_guess=[0.01, 0.005, 0.1, 0.0001, 200])
    randles2 = Randles(initial_guess=[0.01, 0.005, 0.1, 0.0001, 200])
    assert randles1 == randles2

    randles1.fit(f[np.imag(Z) < 0], Z[np.imag(Z) < 0])
    assert randles1 != randles2

    randles2.fit(f[np.imag(Z) < 0], Z[np.imag(Z) < 0])
    assert randles1 == randles2

    randles2.fit(f, Z)
    assert randles1 != randles2


def test_CustomCircuit():
    initial_guess = [0.01, 0.005, 0.1, 0.005, 0.1, 0.001, 200]
    custom_string = "R0-p(R1,C1)-p(R2,C2)-Wo1"
    custom_circuit = CustomCircuit(initial_guess=initial_guess, circuit=custom_string)

    # check get_param_names()
    full_names, all_units = custom_circuit.get_param_names()
    assert full_names == ["R0", "R1", "C1", "R2", "C2", "Wo1_0", "Wo1_1"]
    assert all_units == ["Ohm", "Ohm", "F", "Ohm", "F", "Ohm", "sec"]

    # check _is_fit()
    assert not custom_circuit._is_fit()

    # check predictions from initial_guesses
    high_f = np.array([1e9])
    assert np.isclose(np.real(custom_circuit.predict(high_f)[0]), initial_guess[0])

    # __str()__
    initial_guess = [0.01, 0.005, 0.1]
    custom_string = "R0-p(R1,C1)"
    custom_circuit = CustomCircuit(initial_guess=initial_guess, circuit=custom_string)

    assert (
        str(custom_circuit)
        == "\nCircuit string: R0-p(R1,C1)\n"
        + "Fit: False\n"
        + "\nInitial guesses:\n"
        + "     R0 = 1.00e-02 [Ohm]\n"
        + "     R1 = 5.00e-03 [Ohm]\n"
        + "     C1 = 1.00e-01 [F]\n"
    )

    custom_circuit.fit(f, Z)
    assert custom_circuit._is_fit()
    custom_circuit.plot(f_data=f, Z_data=Z)

    # constants and _ in circuit and no name
    circuit = "R_0-p(R_1,C_1)-Wo_1"
    constants = {"R_0": 0.02, "Wo_1_1": 200}
    initial_guess = [0.005, 0.1, 0.001]
    custom_circuit = CustomCircuit(
        initial_guess=initial_guess, constants=constants, circuit=circuit, name="Test"
    )

    assert (
        str(custom_circuit)
        == "\nName: Test\n"
        + "Circuit string: R_0-p(R_1,C_1)-Wo_1\n"
        + "Fit: False\n"
        + "\nConstants:\n"
        + "    R_0 = 2.00e-02 [Ohm]\n"
        + "  Wo_1_1 = 2.00e+02 [sec]\n"
        + "\nInitial guesses:\n"
        + "    R_1 = 5.00e-03 [Ohm]\n"
        + "    C_1 = 1.00e-01 [F]\n"
        + "  Wo_1_0 = 1.00e-03 [Ohm]\n"
    )

    # incorrect number of initial guesses
    with pytest.raises(ValueError):
        initial_guess = [0.01, 0.005, 0.1, 0.005, 0.1, 0.001, 200]
        custom_string = "R0-p(R1,CPE1)-p(R1,C1)-Wo1"
        custom_circuit = CustomCircuit(
            initial_guess=initial_guess, circuit=custom_string
        )

    # no initial guesses supplied before fitting
    with pytest.raises(ValueError):
        custom_circuit = CustomCircuit()
        custom_circuit.fit(f, Z)

    # incorrect circuit element in circuit
    with pytest.raises(ValueError):
        custom_circuit = CustomCircuit("R0-NotAnElement", initial_guess=[1, 2])

# get example data # __init__() # check initial_guess is loaded in correctly # improper input_guess types raise an TypeError # __eq__() # incorrect comparisons raise a TypeError # fit() # improper data types in fitting raise a TypeError # frequencies not ndarray # frequencies not numeric type # Z not ndarray # Z not complex # mismatched lengths # predict() # improper data types in fitting raise a TypeError # frequencies not ndarray # frequencies not numeric type # plot() # kind = {'nyquist', 'bode'} should return a plt.Axes() object # incorrect kind raises a ValueError # incorrect initial guess length # compare with known fit parameters # compare with known impedance predictions # check altair plotting with a fit circuit # plot() with fitted model # check defaults work if no frequency data is passed # bode plots # missing f_data # nyquist plots # check equality comparisons work # check get_param_names() # check _is_fit() # check predictions from initial_guesses # __str()__ # constants and _ in circuit and no name # incorrect number of initial guesses # no initial guesses supplied before fitting # incorrect circuit element in circuit