import numpy as np
from numpy.linalg import norm

"""
A quadratic data term
"""


class QuadraticDataTerm:
    def __init__(self, opA, b, b_norm, xbound=None, λ=1.0):
        self.opA = opA
        self.b = b
        self.b_norm = b_norm
        self.xbound = xbound
        self.λ = λ

    def apply(self, x):
        v = self.opA.apply(x) - self.b
        return norm(v) ** 2 * (self.λ / 2)

    def diff(self, x):
        v = self.opA.apply(x) - self.b
        j = self.λ * v
        return self.opA.diff_adjdir(j, x)

    def apply_and_diff(self, x):
        v = self.opA.apply(x) - self.b
        val = norm(v) ** 2 * (self.λ / 2)
        j = self.λ * v
        d = self.opA.diff_adjdir(j, x)
        return (val, d)

    def diff_lipschitz_factor(self, xbound=None):
        if hasattr(self.opA, "opnorm"):
            # Linear operator
            return self.λ * self.opA.opnorm() ** 2
        else:
            raise Exception("diff_lipschitz_factor for non nonlinear operators")
            # ‖∇A(x)^*∇F(A(x)) - ∇A(y)^*∇F(A(y))‖
            # = ‖[∇A(x)^*-∇A(y)^*]∇F(A(x)) - ∇A(y)^*[A(y)-A(x)]‖
            # ≤ L_{∇A(x)}(M_A+M_b) + M_{∇A(y)^*} L_A.
            la = self.opA.lipschitz_factor()
            ma = self.opA.bound(xbound=xbound)
            mb = self.b.norm()
            mda = self.opA.diff_bound(xbound=xbound)
            lda = self.opA.diff_adj_lipschitz_factor()
            return self.λ * (lda * (ma + mb) + mda * la)

    # def diff_lipschitz_factor_pair(self, *args):
    #     return self.diff_lipschitz_factor()

    def diff_bound(self, xbound=None):
        return self.λ * (self.opA.codomain_bound(xbound=xbound) + self.b_norm)
