import logging
import pyvisa
import numpy as np
import platform
import time
from tqdm import tqdm as pbar
import matplotlib.pyplot as plt
from qolab.hardware import BasicInstrument
from qolab.data import TraceSetSameX, TraceXY, Trace
from qolab.hardware.rf_generator import QOL_LMX2487
from qolab.hardware.daq import LabJackUE9
from qolab.hardware.power_supply.keysight_e3612a import KeysightE3612A
import qolab.tsdb as tsdb

# this should be done before justpy is called or log formatter does not work
logging.basicConfig(
    format="%(asctime)s %(levelname)8s %(name)s: %(message)s",
    datefmt="%m/%d/%Y %H:%M:%S",
)

l = logging.getLogger("qolab.tsdb")
l.setLevel(logging.INFO)
logger = logging.getLogger("BfieldDriver")
logger.setLevel(logging.INFO)


class BfieldDriver(KeysightE3612A):
    """need to set power supply"""

    def __init__(self, *args, **kwds):
        super().__init__(*args, **kwds)
        self.config[
            "Device type"
        ] = "B field coil driver based on Keysight E3612A power supply"
        self.config["Device model"] = "v0.1"
        self.deviceProperties = self.deviceProperties.union(
            {"B", "Bslope_TperA", "CoilAssignment"}
        )
        """"
        Rough magnetic field calibration of the 3 axes coils in large magnetic shield
        - Ch1: 70mA -> 650 kHz shift for delta m = 2
        - Ch2: 70mA -> 700 kHz shift for delta m = 2
        - Ch2: 70mA -> 659 kHz shift for delta m = 2
        A better calibration obtained on 20220601 see file 
          20220601.magnetic_field_callibration/calibration_currentToB.dat
        """
        # B response to current in a given channel
        # Olivia's Calibration with flux gate magnetometer used from 2023/04/28
        #  +/-     0.0000007
        self._Bslope_TperA = {1: 0.00065859269, 2: 0.00070732580, 3: 0.00066754994}
        """
        # Irina's Calibration used prior 2023/04/27
        self._Bslope_TperA = {
                1: 0.0006571429,
                2: 0.0007085714,
                3: 0.0006675714
                }
        # Eugeniy's Calibration
        self._Bslope_TperA = {
                1: 0.0006574710928926532,
                2: 0.0007064314754023079,
                3: 0.0006635058865577695
                }
        """
        # assuming that Ch1 controls Bz, Ch2 -> Bx, Ch3 -> By
        self._coil_assignment = {"chX": 2, "chY": 3, "chZ": 1}

    def getBslope_TperA(self):
        return self._Bslope_TperA

    def getCoilAssignment(self):
        return self._coil_assignment

    def getB(self):
        Bslope_TperA = self.getBslope_TperA()
        coil_assignment = self.getCoilAssignment()
        chX = coil_assignment["chX"]
        chY = coil_assignment["chY"]
        chZ = coil_assignment["chZ"]

        Ix = self.getChanIout(chX)
        Iy = self.getChanIout(chY)
        Iz = self.getChanIout(chZ)
        Bx = Ix * Bslope_TperA[chX]
        By = Iy * Bslope_TperA[chY]
        Bz = Iz * Bslope_TperA[chZ]
        Bmag = float(np.sqrt(Bx * Bx + By * By + Bz * Bz))
        theta = float(np.arccos(Bz / Bmag))
        phi = float(np.arctan2(By, Bx))
        theta_degree = float(theta / np.pi * 180)
        phi_degree = float(phi / np.pi * 180)
        return {
            "Bmag": Bmag,
            "theta": theta,
            "phi": phi,
            "theta_degree": theta_degree,
            "phi_degree": phi_degree,
            "Bx": float(Bx),
            "By": float(By),
            "Bz": float(Bz),
        }

    def setB(self, Bmag=50e-6, theta=0, phi=0):
        """Sets B field currents based on B (in T) and angles theta, and phi"""
        self._Bmag = Bmag
        self._theta = theta
        self._phi = phi

        Bx = Bmag * np.sin(theta) * np.cos(phi)
        By = Bmag * np.sin(theta) * np.sin(phi)
        Bz = Bmag * np.cos(theta)

        Bslope_TperA = self.getBslope_TperA()
        coil_assignment = self.getCoilAssignment()
        chX = coil_assignment["chX"]
        chY = coil_assignment["chY"]
        chZ = coil_assignment["chZ"]

        Ix = Bx / Bslope_TperA[chX]
        Iy = By / Bslope_TperA[chY]
        Iz = Bz / Bslope_TperA[chZ]

        logger.info(f"Setting {Bmag=}, {theta=} {phi=} in radians")
        logger.info(f"Setting {chX=} to {Ix}")
        logger.info(f"Setting {chY=} to {Iy}")
        logger.info(f"Setting {chZ=} to {Iz}")
        self.setChanIout(chX, Ix)
        self.setChanIout(chY, Iy)
        self.setChanIout(chZ, Iz)
        return Ix, Iy, Iz

    def setBinDegrees(self, Bmag=50e-6, theta=0, phi=0):
        logger.info(f"Setting {Bmag=},  {theta=} {phi=} in degrees")
        return self.setB(Bmag=Bmag, theta=theta / 180 * np.pi, phi=phi / 180 * np.pi)


# TSDB logger setting
tsdb_client = tsdb.Client("influx", "http://qo.physics.wm.edu:8428", database="qolab")
tsdb_ingester = tsdb.Ingester(tsdb_client, batch=11, measurement_prefix="VAMPIRE.VCSEL")

logger.info("Accessing hardware")
rm = pyvisa.ResourceManager()
instr = rm.open_resource("USB0::10893::4354::MY61001869::0::INSTR")
app_nickname = "BfieldDriver"
# Bfield = BfieldDriver(instr, device_nickname='.'.join([app_nickname, 'b_field_driver']), tsdb_ingester=tsdb_ingester)
Bfield = BfieldDriver(
    instr,
    device_nickname=".".join([app_nickname, "b_field_driver"]),
    tsdb_ingester=None,
)


print("Use me as : Bfield.setBinDegrees(theta=Angle, phi=Angle)")