from pydantic import ConfigDict, Field, model_validator
from typing_extensions import Union, Annotated, Literal, List
import matplotlib.pyplot as plt
import numpy as np
from .. import base_model as bm
from ..plot_utils import convert_fig_to_html_img

class ModuleIVV0(bm.ConstructionBase):
    model_config = ConfigDict(json_schema_extra={
        'examples': [
            {
                "module": "MV4002",
                "component_pos": 2,
                "name": "module iv",
                "measurement_date": "2023-01-01T12:00:00+01:00",
                "location": "BU",
                "user_created": "hayden",
                "version": "v0",
                "current":  [0,1,2,3,4,5],
                "voltage":  [0,1,2,3,4,5],
                "k_factor": [0,1,2,3,4,5]
            },
        ],
        'table': 'test',
        'module_types': ["Prototype", "PreProduction", "Production"],
        'component_types': [],
        'description': 'IV curve of a sensor on a module.'
    })
    name: Literal['module iv'] = "module iv"
    version: Literal["v0"] = "v0"

    module: str
    component_pos: int = Field(ge=0, le=3, description="Component position (0-3)")
    
    current: List[float]
    voltage: List[float]
    k_factor: Union[None, List[float]] = None

    @model_validator(mode='after')
    def same_lengths(self):
        if self.current is not None and self.voltage is not None:
            if len(self.current) != len(self.voltage):
                raise ValueError(f'Current and voltage arrays should have the same lengths. Length of Current, Length of Voltage = ({len(self.current)}, {len(self.voltage)})')
        if self.k_factor is not None and self.current is not None:
            if len(self.k_factor) != len(self.current):
                raise ValueError(f'k_factor and current arrays should have the same lengths. Length of k_factor, Length of Current = ({len(self.k_factor)}, {len(self.current)})')
        return self

    def plot(self):
        fig, ax = plt.subplots(figsize=(6, 4))
        ax.plot(self.voltage, self.current, marker='o', label='IV Curve')

        if self.k_factor is None:
            plot_k_factor = [np.nan]
            for i in range(1, len(self.current)):
                delta_v = self.voltage[i] - self.voltage[i - 1]
                current_i = self.current[i]
                current_prev = self.current[i - 1]

                if delta_v == 0 or current_i == 0:
                    plot_k_factor.append(np.nan)
                    continue

                k_val = ((current_i - current_prev) / delta_v) * (self.voltage[i] / current_i)
                plot_k_factor.append(k_val)
        else:
            plot_k_factor = self.k_factor

        ax2 = ax.twinx()
        ax2.plot(self.voltage, plot_k_factor, marker='x', linestyle='--', color='orange', label='k_factor')
        ax2.set_ylabel('k_factor')

        ax.set_xlabel('Voltage (V)')
        ax.set_ylabel('Current (A)')
        ax.set_title(f'Module {self.module} - Component {self.component_pos}')
        ax.grid(True)
        handles1, labels1 = ax.get_legend_handles_labels()
        handles2, labels2 = ax2.get_legend_handles_labels()
        ax.legend(handles1 + handles2, labels1 + labels2, loc='best')
        return fig
    
    def html_display(self):
        fig = self.plot()
        return convert_fig_to_html_img(fig)
    
ModuleIVType = Annotated[Union[ModuleIVV0], Field(discriminator="version")]