import numpy as np
import scipy.sparse

from ..mode_basis import ModeBasis
from .deformable_mirror import DeformableMirror

class SegmentedDeformableMirror(DeformableMirror):
    '''A segmented deformable mirror.

    This deformable mirror class can simulate devices such as those
    made by IrisAO and BMC. All segments are controlled in piston,
    tip and tilt.

    Parameters
    ----------
    segments : ModeBasis
        A mode basis with all segments.
    '''
    def __init__(self, segments):
        self.segments = segments
        self.actuators = np.zeros(len(segments) * 3)
        self.input_grid = segments.grid

    @property
    def segments(self):
        '''The segments of this deformable mirror in a ModeBasis.
        '''
        return self._segments

    @segments.setter
    def segments(self, segments):
        self._segments = segments

        tip = []
        tilt = []

        for segment in segments:
            segment_mean = segment.mean()
            norm = np.mean(segment**2) - segment_mean**2

            tip_mode = segment * segment.grid.x
            if norm != 0:
                beta = ((tip_mode * segment).mean() - segment_mean * tip_mode.mean()) / norm
                tip_mode -= beta * segment

            tip_mode = scipy.sparse.csr_matrix(tip_mode)
            tip_mode.eliminate_zeros()
            tip.append(tip_mode)

            tilt_mode = segment * segment.grid.y
            if norm != 0:
                beta = ((tilt_mode * segment).mean() - segment_mean * tilt_mode.mean()) / norm
                tilt_mode -= beta * segment

            tilt_mode = scipy.sparse.csr_matrix(tilt_mode)
            tilt_mode.eliminate_zeros()
            tilt.append(tilt_mode)

        tip = ModeBasis(tip)
        tilt = ModeBasis(tilt)

        self.influence_functions = segments + tip + tilt

    def get_segment_actuators(self, segment_id):
        '''Get the actuators for an individual segment of the DM.

        Parameters
        ----------
        segment_id : int
            The index of the segment for which to get the actuators.

        Returns
        -------
        piston : scalar
            The piston of the segment in meters.
        tip : scalar
            The tip of the segment in radians.
        tilt : scalar
            The tilt of the segment in radians.
        '''
        piston = self.actuators[segment_id]
        tip = self.actuators[segment_id + len(self._segments)]
        tilt = self.actuators[segment_id + 2 * len(self._segments)]

        return (piston, tip, tilt)

    def set_segment_actuators(self, segment_id, piston, tip, tilt):
        '''Set the actuators for an individual segment of the DM.

        Parameters
        ----------
        segment_id : int
            The index of the segment for which to get the actuators.
        piston : scalar
            The piston of the segment in meters.
        tip : scalar
            The tip of the segment in radians.
        tilt : scalar
            The tilt of the segment in radians.
        '''
        self.actuators[segment_id] = piston
        self.actuators[segment_id + len(self._segments)] = tip
        self.actuators[segment_id + 2 * len(self._segments)] = tilt
