from finesse.simulations.basematrix cimport MatrixSystemSolver, CarrierSignalMatrixSimulation

import numpy as np
cimport numpy as np

ctypedef (double*, double*) ptr_tuple_2

cdef class OpticalBandpassValues(BaseCValues):
    def __init__(self):
        cdef ptr_tuple_2 ptr = (&self.fc, &self.bandwidth)
        cdef tuple params = ("fc", "bandwidth")
        self.setup(params, sizeof(ptr), <double**>&ptr)


cdef class OpticalBandpassConnections:
    """Contains C accessible references to submatrices for
    optical connections for this element.
    """
    def __cinit__(self, object mirror, MatrixSystemSolver mtx):
        # Only 1D arrays of submatrices as no frequency coupling happening
        cdef int Nf = mtx.optical_frequencies.size
        self.P1i_P2o = SubCCSView1DArray(Nf)
        self.P2i_P1o = SubCCSView1DArray(Nf)
        self.ptrs.P1i_P2o = <PyObject**>self.P1i_P2o.views
        self.ptrs.P2i_P1o = <PyObject**>self.P2i_P1o.views


cdef class OpticalBandpassWorkspace(KnmConnectorWorkspace):
    def __init__(self, owner, CarrierSignalMatrixSimulation sim, bint refill):
        super().__init__(
            owner,
            sim,
            refill,
            refill,
            OpticalBandpassConnections(owner, sim.carrier),
            OpticalBandpassConnections(owner, sim.signal) if sim.signal else None,
            OpticalBandpassValues()
        )
        # Casting python objects to known types for faster access
        self.cvalues = self.values
        self.carrier_opt_conns = self.carrier.connections
        if sim.signal:
            self.signal_opt_conns = self.signal.connections
        else:
            self.signal_opt_conns = None


cdef void fill_optical_matrix(
        OpticalBandpassWorkspace ws,
        MatrixSystemSolver matrix,
        optical_bandpass_connections *connections
    ):
    cdef double B = 2*np.pi*ws.cvalues.bandwidth
    cdef complex_t H
    cdef frequency_info_t *frequencies = matrix.optical_frequencies.frequency_info
    cdef Py_ssize_t i

    for i in range(matrix.optical_frequencies.size):
        H = 1 / (1 + 1j*(2*np.pi*abs(frequencies[i].f-ws.cvalues.fc))/B)
        if connections.P1i_P2o:
            (<SubCCSView>connections.P1i_P2o[frequencies[i].index]).fill_negative_za_zm_2(
                H, ws.K12.mtx.ptr, ws.K12.mtx.stride1, ws.K12.mtx.stride2
            )

        if connections.P2i_P1o:
            (<SubCCSView>connections.P2i_P1o[frequencies[i].index]).fill_negative_za_zm_2(
                H, ws.K21.mtx.ptr, ws.K21.mtx.stride1, ws.K21.mtx.stride2
            )


optical_bandpass_carrier_fill = FillFuncWrapper.make_from_ptr(c_fill_carrier)
cdef object c_fill_carrier(ConnectorWorkspace ws):
    fill_optical_matrix(
        ws,
        ws.sim.carrier,
        &(<OpticalBandpassWorkspace>ws).carrier_opt_conns.ptrs
    )

optical_bandpass_signal_fill = FillFuncWrapper.make_from_ptr(c_fill_signal)
cdef object c_fill_signal(ConnectorWorkspace ws):
    fill_optical_matrix(
        ws,
        ws.sim.signal,
        &(<OpticalBandpassWorkspace>ws).signal_opt_conns.ptrs
    )