# -*- coding: utf-8 -*-
"""
Vertical electrical sounding (VES) manager class.
"""
import numpy as np

import pygimli as pg
# from pygimli.frameworks import Modelling, Block1DModelling
from pygimli.frameworks import Block1DModelling, MethodManager1d
from pygimli.frameworks.modelling import DEFAULT_STYLES


class VESModelling(Block1DModelling):
    """Vertical Electrical Sounding (VES) forward operator.

    Attributes
    ----------
    am :
        Part of data basis. Distances between A and M electrodes.
        A is first power, M is first potential electrode.
    bm :
        Part of data basis. Distances between B and M electrodes.
        B is second power, M is first potential electrode.
    an :
        Part of data basis. Distances between A and N electrodes.
        A is first power, N is second potential electrode.
    bn :
        Part of data basis. Distances between B and N electrodes.
        B is second power, N is second potential electrode.
    ab2 :
        Half distance between A and B.
    mn2 :
        Half distance between A and B.
        Only used for input (feeding am etc.).
    """

    def __init__(self, ab2=None, mn2=None, **kwargs):
        r"""Constructor
        """
        self.am = None
        self.bm = None
        self.an = None
        self.bn = None
        self.ab2 = None
        self.mn2 = None

        super(VESModelling, self).__init__(**kwargs)

        if 'dataContainerERT' in kwargs:
            data = kwargs['dataContainerERT']
            if isinstance(data, pg.DataContainerERT):
                kwargs['am'] = [data.sensorPosition(data('a')[i]).distance(
                    data('m')[i]) for i in range(data.size())]