# Copyright 2024 The Pigweed Authors
#
# Licensed under the Apache License, Version 2.0 (the "License"); you may not
# use this file except in compliance with the License. You may obtain a copy of
# the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
# License for the specific language governing permissions and limitations under
# the License.
"""ADC related math functions."""

import logging

_CSV_LOG = logging.getLogger('gonk.csv')

# CHANNEL_MASK should be set to same value used by
#   SelectContinuousReadAdcs in the MCU
CHANNEL_MASK = 0b011100111001
ADC_COUNT = CHANNEL_MASK.bit_count()

CHANNEL_NAMES_ALL = [
    "VDD_EE",
    "VCC18",
    "VDDIOAO18",
    "VDDCPU_B",
    "DCIN",
    "VDDQ(DDR)",
    "VCC33",
    "VCC5V",
    "VDDCPU_A",
    "VSYS3V3",
    "EMMC18",
]

SHUNT_RESISTANCES = [
    0.03,
    0.50,
    0.50,
    0.10,
    0.05,
    0.08,
    0.13,
    0.05,
    0.027,
    0.03,
    0.5,
]

# From INA229 Datasheet
VBUS_VOLTS_PER_COUNT = 195.3125e-6
VSHUNT_VOLTS_PER_COUNT = 312.5e-9


def get_channel_names(channel_mask: int = CHANNEL_MASK) -> list[str]:
    """Returns list of channel names based on channel_mask."""
    retval = []

    for name in CHANNEL_NAMES_ALL:
        if channel_mask & 0x01:
            retval.append(name)
        channel_mask = channel_mask >> 1
    return retval


def log_csv_header(channel_mask: int = CHANNEL_MASK) -> None:
    csv_header = [
        'host_time',
        'delta_micros',
    ]
    for measurement_prefix in ['V', 'I', 'P']:
        csv_header.extend(
            [
                f'{measurement_prefix}_{channel}'
                for channel in get_channel_names(channel_mask)
            ]
        )
    csv_header.append('comment')
    _CSV_LOG.info(', '.join(csv_header))


def get_resistances(channel_mask: int = CHANNEL_MASK) -> list[float]:
    """Returns list of shunt resistance values for the activated channels."""
    retval = []

    for ohms in SHUNT_RESISTANCES:
        if channel_mask & 0x01:
            retval.append(ohms)
        channel_mask = channel_mask >> 1
    return retval


def get_shunt_currents(
    vshunt_measurements, channel_mask: int = CHANNEL_MASK
) -> list[float]:
    """Calculate shunt currents for activated channels."""
    retvals = []
    resistances = get_resistances(channel_mask)
    for i, resistance in enumerate(resistances):
        retvals.append(
            VSHUNT_VOLTS_PER_COUNT * vshunt_measurements[i] / resistance
        )
    return retvals


def get_bus_voltages(vbus_measurements) -> list[float]:
    """Calculate bus voltages from raw vbus measurements."""
    retvals = []
    for measurement in vbus_measurements:
        retvals.append(VBUS_VOLTS_PER_COUNT * measurement)
    return retvals


def calc_power(voltages, currents) -> list[float]:
    """Calculate power from voltage and current measurements."""
    retvals = []
    for i, volts in enumerate(voltages):
        retvals.append(volts * currents[i])
    return retvals