# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# # Coordinate Descent

# ### Lower Bound, Take 4
#
# Ensure feasibility of "Lower Bound, Take 2" by adjusting alpha as necessary.

# + [markdown] heading_collapsed=true
# ### Lower Bound, Take 3
#
# Ensure feasibility by allowing a stochastic mixture with the MLE.  Doesn't work (not DCP).

# + [markdown] hidden=true
# Assume $r_{\min} = 0$ for simplicity.  Idea for online solving $$
# \begin{aligned}
# &\!\min_{Q \succeq 0} &\qquad& \sum_{(w,r)} w r Q_{w, r}, \\
# &\text{subject to} &  & \sum_{(w,r)} w Q_{w,r} = 1, \\
# &                  &  & \sum_{(w,r)} Q_{w,r} = 1, \\
# &                  &  & \sum_n \log(Q_{w_n, r_n}) \geq \phi
# \end{aligned}
# $$ where $\phi = -\frac{1}{2} \chi^{2,\alpha}_{(1)} + \sum_n \log(Q^{\text{(mle)}}_{w_n, r_n})$.  Because the support of $Q$ is at most the empirical support plus $(w_{\min}, 0)$ and $(w_{\max}, 0)$ we will maintain two variables $q_{\min}$ and $q_{\max}$ corresponding to $w_{\min}$ and $w_{\max}$ respectively.  Otherwise we need one primal variable for each data point.  However we will use two primal variables corresponding to $r = 0$ and $r = 1$ respectively.  We will split each datapoint into two points whose fractional counts are proportional to $r$ and $(1 - r)$ respectively.
#
# At time $t$ we receive $(w_t, r_t)$ and we want to determine $q_t$.  We are allowed to choose the new $q_{\min}$ and $q_{\max}$ arbitrarily.  For points $q_{<t}$ we are allowed to scale them by $\psi_0$ and $\psi_1$, corresponding to points with $r = 0$ and $r = 1$ respectively.  We are also allowed to stochastically mix in the maximum likelihood solution according to $\psi_{\text{mle}}$.  We assume $(q_{0,<t}, q_{1,<t}, q_{<t,\min}, q_{<t,\max})$ is feasible before receiving $(w_t, r_t)$.  Then 
#
# $$
# \begin{aligned}
# &\!\min_{q_{0,t}, q_{1,t}, q_{\min}, q_{\max}, \psi_0, \psi_1 \succeq 0, \psi_{\text{mle}} \in [0, 1]} &\qquad& q_{1,t} w_t + \psi_1 v_{<t} + \psi_{\text{mle}} v_{\text{mle}} \\
# &\text{subject to} &  & w_t q_{0,t} + w_t q_{1,t} + w_{\min} q_{\min} + w_{\max} q_{\max} + \psi_0 w_{0,<t} q_{0,<t} + \psi_1 w_{1,<t} q_{1,<t} = 1, \\
# &                  &  & q_{0,t} + q_{1,t} + q_{\min} + q_{\max} + \psi_0 q_{0,<t} + \psi_1 q_{1,<t} = 1, \\
# &                  &  & (1 - r_t) \log(q_{0,t}) + r_t \log(q_{1,t}) + (t - r_{<t}) \log(\psi_0) + r_{<t} \log(\psi_1) \geq \phi - \mathcal{L}_{<t} = -\frac{1}{2} \chi^{2,\alpha}_{(1)} + \mathcal{L}^{\text{(mle)}}_t - \mathcal{L}_{<t}
# \end{aligned}
# $$ where $v_{<t}$ is the previous lower bound, $\mathcal{L}^{\text{(mle)}}_t$ is the mle likelihood of the observed data including point $t$, and $\mathcal{L}_{<t}$ is the previously obtained likelihood of the observed data for the lower bound.

# + [markdown] heading_collapsed=true
# ### Lower Bound, Take 2
#
# Better, but still has feasibility issues as the stream progresses.

# + [markdown] hidden=true
# Assume $r_{\min} = 0$ for simplicity.  Idea for online solving $$
# \begin{aligned}
# &\!\min_{Q \succeq 0} &\qquad& \sum_{(w,r)} w r Q_{w, r}, \\
# &\text{subject to} &  & \sum_{(w,r)} w Q_{w,r} = 1, \\
# &                  &  & \sum_{(w,r)} Q_{w,r} = 1, \\
# &                  &  & \sum_n \log(Q_{w_n, r_n}) \geq \phi
# \end{aligned}
# $$ where $\phi = -\frac{1}{2} \chi^{2,\alpha}_{(1)} + \sum_n \log(Q^{\text{(mle)}}_{w_n, r_n})$.  Because the support of $Q$ is at most the empirical support plus $(w_{\min}, 0)$ and $(w_{\max}, 0)$ we will maintain two variables $q_{\min}$ and $q_{\max}$ corresponding to $w_{\min}$ and $w_{\max}$ respectively.  Otherwise we need one primal variable for each data point.  
#
# At time $t$ we receive $(w_t, r_t)$ and we want to determine $q_t$.  We are allowed to choose the new $q_{\min}$ and $q_{\max}$ arbitrarily.  For points $q_{<t}$ we are allowed to scale them by $\psi_0$ and $\psi_1$, corresponding to points with $r = 0$ and $r = 1$ respectively.  We assume $(q_{0,<t}, q_{1,<t}, q_{<t,\min}, q_{<t,\max})$ is feasible before receiving $(w_t, r_t)$.  Then 
#
# $$
# \begin{aligned}
# &\!\min_{q_t, q_{\min}, q_{\max},\psi_0, \psi_1 \succeq 0} &\qquad& q_t w_t r_t + \psi_1 v_{<t} \\
# &\text{subject to} &  & w_t q_t + w_{\min} q_{\min} + w_{\max} q_{\max} + \psi_0 w_{0,<t} q_{0,<t} + \psi_1 w_{1,<t} q_{1,<t} = 1, \\
# &                  &  & q_t + q_{\min} + q_{\max} + \psi_0 q_{0,<t} + \psi_1 q_{1,<t} = 1, \\
# &                  &  & \log(q_t) + (t - r_{<t}) \log(\psi_0) + r_{<t} \log(\psi_1) \geq \phi - \mathcal{L}_{<t} = -\frac{1}{2} \chi^{2,\alpha}_{(1)} + \mathcal{L}^{\text{(mle)}}_t - \mathcal{L}_{<t}
# \end{aligned}
# $$ where $v_{<t}$ is the previous lower bound, $\mathcal{L}^{\text{(mle)}}_t$ is the mle likelihood of the observed data including point $t$, and $\mathcal{L}_{<t}$ is the previously obtained likelihood of the observed data for the lower bound.
#
# After computing $q_t$ we increment $q_{0,<t}$ and $q_{1,<t}$ proportional to $1 - r_t$ and $r_t$ respectively; and $w_{0,<t}$ and $w_{1,<t}$ by $w_t (1 - r_t)$ and $w_t r_t$ respectively.

# + [markdown] heading_collapsed=true
# ### Lower Bound
#
# Starts out good, but then runs into infeasibility.

# + [markdown] hidden=true
# Assume $r_{\min} = 0$ for simplicity.  Idea for online solving $$
# \begin{aligned}
# &\!\min_{Q \succeq 0} &\qquad& \sum_{(w,r)} w r Q_{w, r}, \\
# &\text{subject to} &  & \sum_{(w,r)} w Q_{w,r} = 1, \\
# &                  &  & \sum_{(w,r)} Q_{w,r} = 1, \\
# &                  &  & \sum_n \log(Q_{w_n, r_n}) \geq \phi
# \end{aligned}
# $$ where $\phi = -\frac{1}{2} \chi^{2,\alpha}_{(1)} + \sum_n \log(Q^{\text{(mle)}}_{w_n, r_n})$.  Because the support of $Q$ is at most the empirical support plus $(w_{\min}, 0)$ and $(w_{\max}, 0)$ we will maintain two variables $q_{\min}$ and $q_{\max}$ corresponding to $w_{\min}$ and $w_{\max}$ respectively.  Otherwise we need one primal variable for each data point.
#
# At time $t$ we receive $(w_t, r_t)$ and we want to determine $q_t$.  We are allowed to choose the new $q_{\min}$ and $q_{\max}$ arbitrarily.  For points $q_{<t}$ we only allowed to scale them by $\psi$.  We assume $(q_{<t}, q_{<t,\min}, q_{<t,\max})$ is feasible before receiving $(w_t, r_t)$.  Then 
#
# $$
# \begin{aligned}
# &\!\min_{q_t, q_{\min}, q_{\max}, \psi \succeq 0} &\qquad& q_t w_t r_t + \psi v_{<t} \\
# &\text{subject to} &  & w_t q_t + w_{\min} q_{\min} + w_{\max} q_{\max} + \psi w_{<t} q_{<t} = 1, \\
# &                  &  & q_t + q_{\min} + q_{\max} + \psi q_{<t} = 1, \\
# &                  &  & \log(q_t) + t \log(\psi) \geq \phi - \mathcal{L}_{<t} = -\frac{1}{2} \chi^{2,\alpha}_{(1)} + \mathcal{L}^{\text{(mle)}}_t - \mathcal{L}_{<t}
# \end{aligned}
# $$ where $v_{<t}$ is the previous lower bound, $\mathcal{L}^{\text{(mle)}}_t$ is the mle likelihood of the observed data including point $t$, and $\mathcal{L}_{<t}$ is the previously obtained likelihood of the observed data for the lower bound. Substituting $q_{<t} + q_{<t,\min} + q_{<t,\max} = 1$ and $w_{<t} q_{<t} + w_{\min} q_{<t,\min} + w_{\max} q_{<t,\max} = 1$ yields 
#
# $$
# \begin{aligned}
# &\!\min_{q_t, q_{\min}, q_{\max} \succeq 0, \psi \in [0, 1]} &\qquad& q_t w_t r_t + \psi v_{<t} \\
# &\text{subject to} &  & w_t q_t + w_{\min} q_{\min} + w_{\max} q_{\max} + \psi (1 - w_{\min} q_{<t,\min} - w_{\max} q_{<t,\max}) = 1, \\
# &                  &  & q_t + q_{\min} + q_{\max} + \psi (1 - q_{<t,\max} - q_{<t,\min}) = 1 \\
# &                  &  & -\log(q_t) - t \log(\psi) - \frac{1}{2} \chi^{2,\alpha}_{(1)} + \mathcal{L}^{\text{(mle)}}_t  - \mathcal{L}_{<t} \leq 0
# \end{aligned}
# $$
#

# + [markdown] heading_collapsed=true
# ### MLE

# + [markdown] hidden=true
# Idea for online solving $$
# \begin{aligned}
# &\!\max_{Q \succeq 0} &\qquad& \sum_n \log(Q_{w_n, r_n}), \\
# &\text{subject to} &  & \sum_{(w,r)} w Q_{w,r} = 1, \\
# &                  &  & \sum_{(w,r)} Q_{w,r} = 1. 
# \end{aligned}
# $$
# Because the support of $Q$ is at most the empirical support plus $w_{\min}$ and $w_{\max}$ we will maintain two variables $q_{\min}$ and $q_{\max}$ corresponding to $w_{\min}$ and $w_{\max}$ respectively.  Otherwise we need one primal variable for each data point.
#
# At time $t$ we receive $(w_t, r_t)$ and we want to determine $q_t$.  We are allowed to choose the new $q_{\min}$ and $q_{\max}$ arbitrarily.  For points $q_{<t}$ we only allowed to scale them by $\psi$.  We assume $(q_{<t}, q_{<t,\min}, q_{<t,\max})$ is feasible before receiving $(w_t, r_t)$.  Then 
#
# $$
# \begin{aligned}
# &\!\max_{q_t, q_{\min}, q_{\max}, \psi > 0} &\qquad& t \log(\psi) + \log(q_t), \\
# &\text{subject to} &  & w_t q_t + w_{\min} q_{\min} + w_{\max} q_{\max} + \psi w_{<t} q_{<t} = 1, \\
# &                  &  & q_t + q_{\min} + q_{\max} + \psi q_{<t} = 1
# \end{aligned}
# $$ 
#
# Substituting $q_{<t} + q_{<t,\min} + q_{<t,\max} = 1$ and $w_{<t} q_{<t} + w_{\min} q_{<t,\min} + w_{\max} q_{<t,\max} = 1$ yields 
#
# $$
# \begin{aligned}
# &\!\max_{q_t, q_{\min}, q_{\max}, \psi > 0} &\qquad& t \log(\psi) + \log(q_t), \\
# &\text{subject to} &  & w_t q_t + w_{\min} q_{\min} + w_{\max} q_{\max} + \psi (1 - w_{\min} q_{<t,\min} - w_{\max} q_{<t,\max}) = 1, \\
# &                  &  & q_t + q_{\min} + q_{\max} + \psi (1 - q_{<t,\max} - q_{<t,\min}) = 1 
# \end{aligned}
# $$
#
# At the beginning of time we can initialize with $$
# \begin{aligned}
# q_{0,\min} &= \frac{1 - w_{\min}}{w_{\max} - w_{\min}} \\
# q_{0,\max} &= \frac{w_{\max} - 1}{w_{\max} - w_{\min}}
# \end{aligned}
# $$
# -

# ### Code

# + code_folding=[0, 1, 26, 79, 177, 247, 315, 432, 517, 660]
class OnlineCoordinateDescentMLE:
    def __init__(self, wmin, wmax):
        from cvxopt import matrix

        assert wmax > 1
        assert wmin >= 0
        assert wmin < wmax
        
        self.wmin = wmin
        self.wmax = wmax
        self.qmin = (wmax - 1) / (wmax - wmin)
        self.qmax = (1 - wmin) / (wmax - wmin)
        self.obj = 0
        self.vmin = 0
        self.lastphi = 0
        
        self.G = matrix([ [ -1, 0, 0, 0 ],
                          [ 0, -1, 0, 0 ],
                          [ 0, 0, -1, 0 ],
                          [ 0, 0, 0, -1 ],
                        ],
                        tc='d').T
        self.h = matrix([ 0, 0, 0, 0 ], tc='d')
        self.b = matrix([ 1 / wmax, 1 ], tc='d')        
        self.t = 0
        
    def update(self, c, w, r):
        from cvxopt import matrix, solvers
        
        assert c > 0
    
        safet = max(self.t, 1)
        x0 = matrix([ c / (c + safet), 
                      self.qmin * safet / (c + safet), 
                      self.qmax * safet / (c + safet), 
                      safet / (c + safet) ], 
                    tc='d')

        def F(x=None, z=None):
            import math
            
            if x is None: return 0, x0
            
            if x[0] <= 0 or x[3] <= 0:
                return None
            
            f = -c * math.log(x[0]) / safet - self.t * math.log(x[3]) / safet
            jf = matrix([ -c / (safet * x[0]), 0, 0, -self.t / (safet * x[3]) ], tc='d').T
            if z is None: return f, jf
            hf = z[0] * matrix([ [ (c / safet) * 1/x[0]**2, 0, 0, 0 ], 
                                 [ 0, 0, 0, 0 ],
                                 [ 0, 0, 0, 0 ],
                                 [ 0, 0, 0, (self.t / safet) * 1/x[3]**2 ] 
                               ], tc='d')
            return f, jf, hf
          
        A = matrix([ 
                     [ float(w) / self.wmax, 
                       self.wmin / self.wmax, 
                      1, 
                      (1 / self.wmax - (self.wmin / self.wmax) * self.qmin - self.qmax) ],
                     [ 1, 1, 1, (1 - self.qmin - self.qmax) ]
                   ],
                   tc='d')
                
        soln = solvers.cp(F=F, G=self.G, h=self.h, A=A.T, b=self.b, options={'show_progress': False})
        from pprint import pformat
        assert soln['status'] == 'optimal', pformat([ soln, self.t ])
        
        self.obj -= safet * soln['primal objective']
        self.lastq = soln['x'][0]
        self.qmin = soln['x'][1]
        self.qmax = soln['x'][2]
        self.lastphi = soln['x'][3]
        self.vmin = soln['x'][0] * w * r + soln['x'][3] * self.vmin
        self.t += c
        
        return self.lastq

class OnlineCoordinateDescentLB: 
    class Flass:
        def __init__(self):
            pass
        
    def __init__(self, wmin, wmax, alpha): 
        assert wmax > 1
        assert wmin >= 0
        assert wmin < wmax
        
        self.wmin = wmin
        self.wmax = wmax
        self.qmin = (wmax - 1) / (wmax - wmin)
        self.qmax = (1 - wmin) / (wmax - wmin)
        
        self.vlb = 0
        self.wq0 = 0
        self.wq1 = 0
        self.q0t = 0
        self.q1t = 0
        self.t = 0
        self.rt = 0
        self.llb = 0
        
        self.alpha = alpha
        self.mle = OnlineCoordinateDescentMLE(wmin=wmin, wmax=wmax)
        
        from scipy.stats import chi2
        import cvxpy as cp
        
        self.vars = OnlineCoordinateDescentLB.Flass()
        self.vars.qt = cp.Variable(nonneg=True)
        self.vars.qmin = cp.Variable(nonneg=True)
        self.vars.qmax = cp.Variable(nonneg=True)
        self.vars.psi0 = cp.Variable(nonneg=True)
        self.vars.psi1 = cp.Variable(nonneg=True)
        
        self.params = OnlineCoordinateDescentLB.Flass()
        self.params.w = cp.Parameter(nonneg=True)
        self.params.wr = cp.Parameter(nonneg=True)
        self.params.vlb = cp.Parameter(nonneg=True)
        self.params.wq0 = cp.Parameter(nonneg=True)
        self.params.wq1 = cp.Parameter(nonneg=True)
        self.params.q0t = cp.Parameter(nonneg=True)
        self.params.q1t = cp.Parameter(nonneg=True)
        self.params.c = cp.Parameter(nonneg=True)
        self.params.tminusrt = cp.Parameter(nonneg=True)
        self.params.rt = cp.Parameter(nonneg=True)
        self.params.constraintrhs = cp.Parameter()
        
        self.prob = cp.Problem(cp.Minimize(self.params.wr * self.vars.qt + self.params.vlb * self.vars.psi1), [
              self.params.w * self.vars.qt
            + (self.wmin / self.wmax) * self.vars.qmin 
            + self.vars.qmax 
            + self.params.wq0 * self.vars.psi0 
            + self.params.wq1 * self.vars.psi1
            == 1 / self.wmax,  
              self.vars.qt 
            + self.vars.qmin 
            + self.vars.qmax 
            + self.params.q0t * self.vars.psi0
            + self.params.q1t * self.vars.psi1 == 1,
              self.params.c * cp.log(self.vars.qt) 
            + self.params.tminusrt * cp.log(self.vars.psi0)
            + self.params.rt * cp.log(self.vars.psi1) 
            >= self.params.constraintrhs
        ])
        
    def innersolve(self, c, w, r, alpha):
        from scipy.stats import chi2

        safet = max(self.t, 1)
        halfchisq = 0.5 * chi2.isf(q=alpha, df=1)
        
        self.params.w.value = w / self.wmax
        self.params.wr.value = w * r
        self.params.vlb.value = self.vlb
        self.params.wq0.value = self.wq0 / self.wmax
        self.params.wq1.value = self.wq1 / self.wmax
        self.params.q0t.value = self.q0t
        self.params.q1t.value = self.q1t
        self.params.c.value = c / safet
        self.params.tminusrt.value = (self.t - self.rt) / safet
        self.params.rt.value = self.rt / safet
        self.params.constraintrhs.value = (-halfchisq + self.mle.obj - self.llb) / safet
        
        self.prob.solve(verbose=False)
        
        return (self.prob.value, 
                self.prob.status, 
                (1 - r) * self.vars.qt.value if self.vars.qt.value is not None else None, 
                r * self.vars.qt.value if self.vars.qt.value is not None else None, 
                self.vars.qmin.value, 
                self.vars.qmax.value, 
                self.vars.psi0.value, 
                self.vars.psi1.value
               )
        
    def updatev3(self, c, w, r):
        import cvxpy as cp
        import math
        from scipy.special import xlogy
        
        qmle = self.mle.update(c, w, r)
        
        q0t = cp.Variable(nonneg=True)
        q1t = cp.Variable(nonneg=True)
        qmin = cp.Variable(nonneg=True)
        qmax = cp.Variable(nonneg=True)
        psi0 = cp.Variable(nonneg=True)
        psi1 = cp.Variable(nonneg=True)
        psimle = cp.Variable(nonneg=True)
        
        safet = max(self.t, 1)

#        prob = cp.Problem(cp.Minimize((float(w) * q1t + self.vlb * psi1) * (1 - psimle) + self.mle.vhat * psimle), [
        prob = cp.Problem(cp.Minimize(float(w) * q1t + self.vlb * psi1 + self.mle.vhat * psimle), [
              float(w / self.wmax) * q0t
            + float(w / self.wmax) * q1t
            + (self.wmin / self.wmax) * qmin 
            + qmax 
            + psi0 * (self.wq0 / self.wmax)
            + psi1 * (self.wq1 / self.wmax) 
            == 1 / self.wmax,  
            q0t + q1t + qmin + qmax + psi0 * self.q0t + psi1 * self.q1t == 1,
              float(c * (1 - r) / safet) * cp.log(q0t) 
            + float(c * r / safet) * cp.log(q1t) 
            + float((self.t - self.rt) / safet) * cp.log(psi0)
            + float(self.rt / safet) * cp.log(psi1) 
#             >= ((-self.halfchisq + self.mle.obj) / safet) * cp.inv_pos(1 - psimle) - (self.llb / safet),
            >= ((-self.halfchisq + self.mle.obj) / safet) * (1 + psimle) - (self.llb / safet),
            psimle <= 1,
            psimle >= 0.9
        ])
        prob.solve(verbose=False)
        assert prob.status[:7] == 'optimal', prob.solve(verbose=True)
        
        self.vlb = (w * q1t.value + self.vlb * psi1.value) * (1 - psimle.value) + psimle.value * self.mle.vhat
        self.qmin = qmin.value
        self.qmax = qmax.value
        self.llb += (  xlogy(c * (1 - r), q0t.value)
                     + xlogy(c * r, q1t.value)
                     + xlogy(self.t - self.rt, psi0.value)
                     + xlogy(self.rt, psi1.value)
                    )
        self.llb = (1 - psimle.value) * self.llb + psimle.value * self.mle.obj
     
        self.t += c
        self.rt += c * r
        self.q0t = q0t.value + psi0.value * self.q0t
        self.wq0 = w * q0t.value + psi0.value * self.wq0
        self.q1t = q1t.value + psi1.value * self.q1t
        self.wq1 = w * q1t.value + psi1.value * self.wq1
        
        myq = (1 - r) * q0t.value + r * q1t.value
        myq = (1 - psimle.value) * myq + psimle.value * qmle
        
        return myq, { 'self.q0t': self.q0t, 
                      'self.q1t': self.q1t,
                      'qmin': qmin.value,
                      'qmax': qmax.value,
                      'vlb': self.vlb,
                      'llb': self.llb,
                      'vmle': self.mle.vhat,
                      'self.rt': self.rt,
                      'psimle': psimle.value,
                    }

    def __initturg__(self, wmin, wmax, alpha): 
        assert wmax > 1
        assert wmin >= 0
        assert wmin < wmax
        
        self.wmin = wmin
        self.wmax = wmax
        self.qmin = (wmax - 1) / (wmax - wmin)
        self.qmax = (1 - wmin) / (wmax - wmin)
        
        self.vlb = 0
        self.wq0 = 0
        self.wq1 = 0
        self.q0t = 0
        self.q1t = 0
        self.t = 0
        self.rt = 0
        self.llb = 0
        
        self.alpha = alpha
        self.mle = OnlineCoordinateDescentMLE(wmin=wmin, wmax=wmax)
        
        from scipy.stats import chi2
        import cvxpy as cp
        
        self.vars = OnlineCoordinateDescentLB.Flass()
        self.vars.q0t = cp.Variable(nonneg=True)
        self.vars.q1t = cp.Variable(nonneg=True)
        self.vars.qmin = cp.Variable(nonneg=True)
        self.vars.qmax = cp.Variable(nonneg=True)
        self.vars.psi0 = cp.Variable(nonneg=True)
        self.vars.psi1 = cp.Variable(nonneg=True)
        
        self.params = OnlineCoordinateDescentLB.Flass()
        self.params.w = cp.Parameter(nonneg=True)
        self.params.wcost = cp.Parameter(nonneg=True)
        self.params.vlb = cp.Parameter(nonneg=True)
        self.params.wq0 = cp.Parameter(nonneg=True)
        self.params.wq1 = cp.Parameter(nonneg=True)
        self.params.q0t = cp.Parameter(nonneg=True)
        self.params.q1t = cp.Parameter(nonneg=True)
        self.params.coneminusr = cp.Parameter(nonneg=True)
        self.params.cr = cp.Parameter(nonneg=True)
        self.params.tminusrt = cp.Parameter(nonneg=True)
        self.params.rt = cp.Parameter(nonneg=True)
        self.params.constraintrhs = cp.Parameter()
        
        self.prob = cp.Problem(cp.Minimize(self.params.wcost * self.vars.q1t + self.params.vlb * self.vars.psi1), [
              self.params.w * self.vars.q0t
            + self.params.w * self.vars.q1t
            + (self.wmin / self.wmax) * self.vars.qmin 
            + self.vars.qmax 
            + self.params.wq0 * self.vars.psi0 
            + self.params.wq1 * self.vars.psi1
            == 1 / self.wmax,  
              self.vars.q0t 
            + self.vars.q1t 
            + self.vars.qmin 
            + self.vars.qmax 
            + self.params.q0t * self.vars.psi0
            + self.params.q1t * self.vars.psi1 == 1,
              self.params.coneminusr * cp.log(self.vars.q0t) 
            + self.params.cr * cp.log(self.vars.q1t) 
            + self.params.tminusrt * cp.log(self.vars.psi0)
            + self.params.rt * cp.log(self.vars.psi1) 
            >= self.params.constraintrhs
        ])

    def innersolveflass(self, c, w, r, alpha):   
        # doesn't work, not sure why (?)

        from cvxopt import matrix, spdiag, solvers
        import numpy as np
        from scipy.stats import chi2
        
        assert 0 < c
        assert 0 <= r
        assert r <= 1

        safet = max(self.t, 1)
        halfchisq = 0.5 * chi2.isf(q=alpha, df=1)
         
        print([
            w / self.wmax,
            w,
            self.vlb,
            self.wq0 / self.wmax,
            self.wq1 / self.wmax,
            self.q0t,
            self.q1t,
            c * (1 - r) / safet,
            c * r / safet,
            (self.t - self.rt) / safet,
            self.rt / safet,
            (-halfchisq + self.mle.obj - self.llb) / safet
        ])
        
        
        G = matrix(-np.eye(6), tc='d')
        h = matrix(0, size=(6,1), tc='d')
        
        A = matrix([ [ w / self.wmax,
                       w / self.wmax,
                       self.wmin / self.wmax, 
                       1, 
                       self.wq0 / self.wmax,
                       self.wq1 / self.wmax 
                     ],
                     [ 1, 1, 1, 1, self.q0t, self.q1t ]
                   ],
                   tc='d')
        b = matrix([ 1 / self.wmax, 1 ], tc='d')
        
        cost = matrix(0, size=(6,1), tc='d')
        cost[1] = float(w) 
        cost[5] = self.vlb 
        
        x0 = matrix([ c * (1 - r) / (c + safet),
                      c * r / (c + safet),
                      self.qmin * safet / (c + safet), 
                      self.qmax * safet / (c + safet), 
                      safet / (c + safet),
                      safet / (c + safet)
                    ], tc='d')

        def F(x=None, z=None):
            from scipy.special import xlogy

            if x is None: return 1, x0
            
            if any(z < 0 for z in x):
                 return None
                
            f = (  xlogy(c * (1 - r), x[0])
                 + xlogy(c * r, x[1])
                 + xlogy(self.t - self.rt, x[3])
                 + xlogy(self.rt, x[4])
                 + halfchisq
                 - self.mle.obj
                 + self.llb)
            f *= -1 / safet
            
            jf = matrix(0, size=(1, 6), tc='d')
            jf[0] = c * (1 - r) / x[0] if c * (1 - r) > 0 else 0 
            jf[1] = c * r / x[1] if c * r > 0 else 0 
            jf[4] = (self.t - self.rt) / x[4] if self.t > self.rt else 0
            jf[5] = self.rt / x[5] if self.rt > 0 else 0
            jf *= -1 / safet
            
            if z is None: return f, jf
            
            hf = spdiag([
                -c * (1 - r) / x[0]**2 if c * (1 - r) > 0 else 0,
                 -c * r / x[1]**2 if c * r > 0 else 0,               
                0,
                0,
                -(self.t - self.rt) / x[4]**2 if self.t > self.rt else 0,
                -self.rt / x[5]**2 if self.rt > 0 else 0 
            ]) 

            hf *= -z[0] / safet

            return f, jf, hf

        soln = solvers.cpl(c=cost, F=F, G=G, h=h, A=A.T, b=b)

        from pprint import pformat
        import numpy
        assert soln['status'][:7] == 'optimal', pformat({ 'soln': soln, 
                                                  'solnx': [ z for z in soln['x'] ], 
                                                  'datum': (c, w, r),
                                                  'F(x=x0)': F(x=x0),
                                                  'A': numpy.matrix(A),
                                                  'b': [ z for z in b ],
                                                  'A.x0 - b': [ z for z in A.T*x0 - b ],
                                                  'G.x0 - h': [ z for z in G*x0 - h ],
                                                  'F(x=soln)': F(x=soln['x']),
                                                  'A.x - b': [ z for z in A.T*soln['x'] - b ],
                                                  'G.x - h': [ z for z in G*soln['x'] - h ],
                                                })
        return (soln['primal objective'], soln['status'], 
                (1 - r) * soln['x'][0],
                r * soln['x'][0]
               ) + tuple(soln['x'][1:])

    def innersolveturg(self, c, w, r, alpha):
        from scipy.stats import chi2

        safet = max(self.t, 1)
        halfchisq = 0.5 * chi2.isf(q=alpha, df=1)
        
        self.params.w.value = w / self.wmax
        self.params.wcost.value = w
        self.params.vlb.value = self.vlb
        self.params.wq0.value = self.wq0 / self.wmax
        self.params.wq1.value = self.wq1 / self.wmax
        self.params.q0t.value = self.q0t
        self.params.q1t.value = self.q1t
        self.params.coneminusr.value = c * (1 - r) / safet
        self.params.cr.value = c * r / safet
        self.params.tminusrt.value = (self.t - self.rt) / safet
        self.params.rt.value = self.rt / safet
        self.params.constraintrhs.value = (-halfchisq + self.mle.obj - self.llb) / safet
        
        self.prob.solve(verbose=False)
        
        return (self.prob.value, 
                self.prob.status, 
                self.vars.q0t.value, 
                self.vars.q1t.value, 
                self.vars.qmin.value, 
                self.vars.qmax.value, 
                self.vars.psi0.value, 
                self.vars.psi1.value
               )
                
    def update(self, c, w, r):
        import math
        from scipy.special import xlogy
        
        self.mle.update(c, w, r)
        
        alpha = self.alpha
        
        (pvalue, pstatus, q0t, q1t, qmin, qmax, psi0, psi1) = self.innersolve(c, w, r, alpha)
        
        if pstatus[:7] != 'optimal':
            alphalb = 0
            alphaub = alpha
            
            while alphaub - alphalb >= 1e-3:
                alphatest = 0.5 * (alphalb + alphaub)
                (pvalue, pstatus, q0t, q1t, qmin, qmax, psi0, psi1) = self.innersolve(c, w, r, alphatest)
            
                if pstatus[:7] == 'optimal':
                    alphalb = alphatest
                else:
                    alphaub = alphatest
                    
            alpha = alphalb
            (pvalue, pstatus, q0t, q1t, qmin, qmax, psi0, psi1) = self.innersolve(c, w, r, alpha)
            
        assert pstatus[:7] == 'optimal', { 'alpha': alpha, 'pstatus': pstatus }
        
        self.vlb = pvalue
        self.qmin = qmin
        self.qmax = qmax
        self.llb += (  xlogy(c * (1 - r), q0t)
                     + xlogy(c * r, q1t)
                     + xlogy(self.t - self.rt, psi0)
                     + xlogy(self.rt, psi1)
                    )
     
        self.t += c
        self.rt += c * r
        self.q0t = q0t + psi0 * self.q0t
        self.wq0 = w * q0t + psi0 * self.wq0
        self.q1t = q1t + psi1 * self.q1t
        self.wq1 = w * q1t + psi1 * self.wq1
        
        return (1 - r) * q0t + r * q1t, { # 'self.q0t': self.q0t, 
#                                           'self.q1t': self.q1t,
#                                           'qmin': qmin.value,
#                                           'qmax': qmax.value,
                                            'vlb': self.vlb,
                                            'vmle': self.mle.vmin,
#                                           'self.rt': self.rt,
                                            'alpha': alpha,
                                        }

    def __initv1__(self, wmin, wmax, alpha):
        from scipy.stats import chi2
        from cvxopt import matrix

        assert wmax > 1
        assert wmin >= 0
        assert wmin < wmax
        
        self.wmin = wmin
        self.wmax = wmax
        self.qmin = (wmax - 1) / (wmax - wmin)
        self.qmax = (1 - wmin) / (wmax - wmin)
        
        self.G = matrix([ [ -1, 0, 0, 0 ],
                          [ 0, -1, 0, 0 ],
                          [ 0, 0, -1, 0 ],
                          [ 0, 0, 0, -1 ],
                        ],
                        tc='d').T
        self.h = matrix([ 0, 0, 0, 0 ], tc='d')
        self.b = matrix([ 1 / wmax, 1 ], tc='d')        
        self.t = 0

        self.halfchisq = 0.5 * chi2.isf(q=alpha, df=1)
        self.llb = 0
        self.vlb = 0
        
        self.mle = OnlineCoordinateDescentMLE(wmin=wmin, wmax=wmax)
    
    def updatev1(self, c, w, r):
        import cvxpy as cp
        import math
        
        self.mle.update(c, w, r)
                
        qt = cp.Variable(nonneg=True)
        qmin = cp.Variable(nonneg=True)
        qmax = cp.Variable(nonneg=True)
        psi = cp.Variable(nonneg=True)
        
        safet = max(self.t, 1)

        prob = cp.Problem(cp.Minimize(float(w * r) * qt + self.vlb * psi), [
              float(w / self.wmax) * qt 
            + (self.wmin / self.wmax) * qmin 
            + qmax 
            + (1 / self.wmax - (self.wmin / self.wmax) * self.qmin - self.qmax) * psi == 1 / self.wmax,
            qt + qmin + qmax + (1 - self.qmin - self.qmax) * psi == 1,
            float(c / safet) * cp.log(qt) + (self.t / safet) * cp.log(psi) >= (-self.halfchisq + self.mle.obj - self.llb) / safet
        ])
        prob.solve(verbose=False)
        
        if prob.status[:7] != 'optimal':
            # just maximize likelihood to recover (?)
            pass
        
        from pprint import pformat
        assert prob.status[:7] == 'optimal', pformat({ 'datum': [ c, w, r ], 
                                                   'x': [ z.value for z in [ qt, qmin, qmax, psi ] ],
                                                   'prob': prob.status,
                                                   't': self.t, 
                                                   'lmle': self.mle.obj, 
                                                   'llb': self.llb,# + c * math.log(qt.value) + self.t * math.log(psi.value),
                                                   'halfchisq': self.halfchisq,
                                                 })
        
        self.vlb = prob.value
        self.qmin = qmin.value
        self.qmax = qmax.value
        self.lastphi = psi.value
        self.llb += c * math.log(qt.value) + self.t * math.log(psi.value)
        self.t += c

        return { 'soln': [ z.value for z in [ qt, qmin, qmax, psi ] ], 'lmle': self.mle.obj, 
                 'llb': self.llb, 'halfchisq': self.halfchisq, 'vlb': self.vlb, 'vhat': self.mle.vhat }
       
        
        if False:
            from cvxopt import matrix, solvers
            import math

            assert c > 0 

            self.mle.update(c, w, r)
            lmle = self.mle.obj

            safet = max(self.t, 1)
            x0 = matrix([ c / (c + safet), 
                          self.qmin * safet / (c + safet), 
                          self.qmax * safet / (c + safet), 
                          safet / (c + safet) ], tc='d')

            def F(x=None, z=None):
                if x is None: return 1, x0

                if x[0] <= 0 or x[3] <= 0:
                     return None

                f = -c * math.log(x[0]) - self.t * math.log(x[3]) - self.halfchisq + lmle - self.llb
                f /= safet
                jf = matrix([ -c / (safet * x[0]), 0, 0, -self.t / (safet * x[3]) ], tc='d').T
                if z is None: return f, jf
                hf = z[0] * matrix([ [ (c / safet) * 1/x[0]**2, 0, 0, 0 ], 
                                     [ 0, 0, 0, 0 ],
                                     [ 0, 0, 0, 0 ],
                                     [ 0, 0, 0, (self.t / safet) * 1/x[3]**2 ] 
                                   ], tc='d')
                return f, jf, hf

            cost = matrix([ float(w * r), 0, 0, self.vlb ], tc='d')
            A = matrix([ 
                         [ float(w) / self.wmax, 
                           self.wmin / self.wmax, 
                          1, 
                          (1 / self.wmax - (self.wmin / self.wmax) * self.qmin - self.qmax) ],
                         [ 1, 1, 1, (1 - self.qmin - self.qmax) ]
                       ],
                       tc='d')

            soln = solvers.cpl(c=cost, F=F, G=self.G, h=self.h, A=A.T, b=self.b, options={'show_progress': True, 
                                                                                          'maxiters': 100 })
            from pprint import pformat
            assert soln['status'] == 'optimal', pformat({ 'datum': [ c, w, r ], 
                                                          'x': [ z for z in soln['x'] ], 
                                                          'soln': soln,
                                                          't': self.t, 
                                                          'lmle': lmle, 
                                                          'llb': self.llb, 
                                                          'F(x0)': F(x=x0),
                                                          'F(x)': F(soln['x'])
                                                        })
            print(pformat((F(x=x0), F(x=soln['x']))))

            self.vlb = soln['primal objective']
            self.qmin = soln['x'][1]
            self.qmax = soln['x'][2]
            self.lastphi = soln['x'][3]
            self.llb += c * math.log(soln['x'][0]) + self.t * math.log(soln['x'][3])
            self.t += c

            return { 'soln': [ z for z in soln['x'] ], 'lmle': lmle, 
                     'llb': self.llb, 'halfchisq': self.halfchisq, 'vlb': self.vlb, 'vhat': self.mle.vhat }

class Test:
    def flass():
        import environments.ControlledRangeVariance
        import MLE.MLE
        
        def batchtoonline(samples, seed=45):
            import numpy as np

            state = np.random.RandomState(seed)

            n = sum(c for c, w, r in samples)
            while n > 0:
                p = np.array([ c for c, w, r in samples ], dtype='float64') / n
                what = state.choice(len(samples), p=p)
                c = min(samples[what][0], 1)
                yield (c, samples[what][1], samples[what][2])
                samples[what] = (samples[what][0] - c, samples[what][1], samples[what][2]) 
                n -= c
                
        env = environments.ControlledRangeVariance.ControlledRangeVariance(seed=45, wsupport=[0,2,1000], expwsq=100)
        _, samples = env.sample(1 + (1 << 16))
        ocd = OnlineCoordinateDescentLB(wmin=0, wmax=1000, alpha=0.05)

        from pprint import pformat
        
        t = 0
        for n, (c, w, r) in enumerate(batchtoonline(samples)):
            qt, qex = ocd.update(c, w, r)
            t += c
            if (n & (n - 1) == 0 and n & 0xAAAAAAAA == 0) or w == 1000:
                vlb = '{:.5}'.format(qex['vlb'])
                vmle = '{:.5}'.format(qex['vmle'])
                print([ c, w, r, t, '{:.3g}'.format(qt), '{:.3f}'.format(t * qt), { 'vlb': vlb, 'vmle': vmle, 'alpha': qex['alpha'] } ], 
                      flush=True)

Test.flass()

# + [markdown] heading_collapsed=true
# # Histogram Based

# + hidden=true
from importlib import reload
import environments.ControlledRangeVariance
import MLE.MLE 

reload(environments.ControlledRangeVariance)
reload(MLE.MLE)

def batchtoonline(samples, seed=45):
    import numpy as np
    
    state = np.random.RandomState(seed)
    
    n = sum(c for c, w, r in samples)
    while n > 0:
        p = np.array([ c for c, w, r in samples ], dtype='float64') / n
        what = state.choice(len(samples), p=p)
        c = min(samples[what][0], 1)
        yield (c, samples[what][1], samples[what][2])
        samples[what] = (samples[what][0] - c, samples[what][1], samples[what][2]) 
        n -= c
    
env = environments.ControlledRangeVariance.ControlledRangeVariance(seed=45, wsupport=[0,2,1000], expwsq=100)
happrox = MLE.MLE.Online.HistApprox(wmin=0, wmax=1000, numbuckets=10)
onlineci = MLE.MLE.Online.CI(wmin=0, wmax=1000, rmin=0, rmax=1, alpha=0.05)
onlinemle = MLE.MLE.Online.MLE(wmin=0, wmax=1000, rmin=0, rmax=1)
_, samples = env.sample(1 + (1 << 16))
from pprint import pformat
print(pformat(samples), flush=True)

t = 0
for n, (c, w, r) in enumerate(batchtoonline(samples)):
    happrox.update(c, w, r)
    onlineci.update(happrox.iterator)
    onlinemle.update(happrox.iterator)
    t += c
    
    if (n & (n - 1) == 0 and n & 0xAAAAAAAA == 0) or w == 1000:
        vmle = MLE.MLE.estimate(happrox.iterator, wmin=0, wmax=1000)[1]['vmin']
        vlb = MLE.MLE.asymptoticconfidenceinterval(happrox.iterator, wmin=0, wmax=1000)[0][0]
        print([ c, w, r, t, 'n/a', 'n/a', { 'vlb': '{:.3f}'.format(vlb), 'vmle': '{:.3f}'.format(vmle), 'alpha': 0.05 } ], 
                              flush=True)
        
#         from pprint import pformat
#         print(pformat(
#                 {
#                     'n': n,
#                     'onlineci': onlineci.getqfunc(),
#                     'onlinemle': onlinemle.getqfunc(),
#                     'batchmle': MLE.MLE.estimate(happrox.iterator, wmin=0, wmax=1000), 
#                     'batchci': MLE.MLE.asymptoticconfidenceinterval(happrox.iterator, wmin=0, wmax=1000)
#                 }), 
#               flush=True)
# -

# # Forest Covertype

# + [markdown] heading_collapsed=true
# ### Full Information Online Learning

# + code_folding=[0, 1] hidden=true
class OnlineDRO:
    class EasyAcc:
        def __init__(self):
            self.n = 0
            self.sum = 0
            
        def __iadd__(self, other):
            self.n += 1
            self.sum += other
            return self
            
        def mean(self):
            return self.sum / max(self.n, 1)
 
    def flass():
        from sklearn.datasets import fetch_covtype
        from sklearn.decomposition import PCA
        from sklearn.linear_model import SGDClassifier
        from sklearn.metrics import accuracy_score
        from math import ceil
        import numpy as np
  
        cov = fetch_covtype()
        cov.data = PCA(whiten=True).fit_transform(cov.data)
        classes = np.unique(cov.target - 1)
        ndata = len(cov.target)
        order = np.random.RandomState(seed=42).permutation(ndata)
        ntrain = ceil(0.2 * ndata)
        Object = lambda **kwargs: type("Object", (), kwargs)()
        train = Object(data = cov.data[order[:ntrain]], target = cov.target[order[:ntrain]] - 1)
        test = Object(data = cov.data[order[ntrain:]], target = cov.target[order[ntrain:]] - 1)
        
        for lr in (1, ):
            print("**** lr = {} ****".format(lr))
            print('{:8.8s}\t{:8.8s}\t{:10.10s}'.format('n', 'emp loss', 'since last'))
            
            classweights = { k: lr for k, _ in Counter(train.target).items() }
            cls = SGDClassifier(loss='log', class_weight=classweights, shuffle=False)
            loss = OnlineDRO.EasyAcc()
            sincelast = OnlineDRO.EasyAcc()
            blocksize = 32

            for pno in range(1):
                order = np.random.RandomState(seed=42+pno).permutation(len(train.data))
                for n, ind in enumerate(zip(*(iter(order),)*blocksize)):
                    v = np.array([ np.outer(t, np.append(t, [1])).ravel() for z in ind for t in ( train.data[z], ) ])
                    actual = [ train.target[z] for z in ind ]
                    if n > 0:
                        pred = cls.predict(v)
                        for p, a in zip(pred, actual):
                            loss += 0 if p == a else 1
                            sincelast += 0 if p == a else 1
                        if (n & (n - 1) == 0): # and n & 0xAAAAAAAA == 0):
                            print('{:<8d}\t{:<8.3f}\t{:<10.3f}'.format(loss.n, loss.mean(), sincelast.mean()), flush=True)
                            sincelast = OnlineDRO.EasyAcc()

                    cls.partial_fit(v, actual, classes=classes)

                print('{:<8d}\t{:<8.3f}\t{:<10.3f}'.format(loss.n, loss.mean(), sincelast.mean()), flush=True)
                sincelast = OnlineDRO.EasyAcc()

                preds = cls.predict(np.array([np.outer(d, np.append(d, [1])).ravel() for d in test.data]))
                ascores = []
                for b in range(16):
                    bootie = np.random.RandomState(90210+b).choice(len(test.target), replace=True, size=len(test.target))
                    ascores.append(accuracy_score(y_true=test.target[bootie], y_pred=preds[bootie]))

                print("test accuracy: {}".format(np.quantile(ascores, [0.05, 0.5, 0.95])))
    
OnlineDRO.flass()      


# + [markdown] heading_collapsed=true
# ### Partial Information Online Learning, Softmax Logging Policy
#
# Uniform $(\tau = 0)$ and softmax $(\tau = 4)$ are pretty similar for off-policy learning but uniform has larger regret.

# + code_folding=[0, 1, 54] hidden=true
class OnlineDRO:
    class EasyAcc:
        def __init__(self):
            self.n = 0
            self.sum = 0
            
        def __iadd__(self, other):
            self.n += 1
            self.sum += other
            return self
            
        def mean(self):
            return self.sum / max(self.n, 1)
 
    def flass():
        from scipy.special import softmax
        from sklearn.datasets import fetch_covtype
        from sklearn.decomposition import PCA
        from sklearn.linear_model import SGDClassifier
        from sklearn.metrics import accuracy_score
        from math import ceil
        import numpy as np
        
        cov = fetch_covtype()
        cov.data = PCA(whiten=True).fit_transform(cov.data)
        classes = np.unique(cov.target - 1)
        ndata = len(cov.target)
        order = np.random.RandomState(seed=42).permutation(ndata)
        ntrain = ceil(0.2 * ndata)
        Object = lambda **kwargs: type("Object", (), kwargs)()
        train = Object(data = cov.data[order[:ntrain]], target = cov.target[order[:ntrain]] - 1)
        test = Object(data = cov.data[order[ntrain:]], target = cov.target[order[ntrain:]] - 1)
        
        blocksize = 32
        for lr, tau in ( (x, y) for x in np.logspace(-2.5, -2, 1) for y in (0, 4, ) ):
            print("*** lr = {} tau = {} ***".format(lr, tau), flush=True)
            print('{:8.8s}\t{:8.8s}\t{:10.10s}\t{:8.8s}\t{:10.10s}'.format(
                'n', 'emp loss', 'since last', 'log pv', 'since last')
            )
            
            cls = SGDClassifier(loss='log', shuffle=False)
            loss = OnlineDRO.EasyAcc()
            sincelast = OnlineDRO.EasyAcc()
            logpv = OnlineDRO.EasyAcc()
            logpvsl = OnlineDRO.EasyAcc()
            
            loggerrand = np.random.RandomState(seed=2112)
            logchoices = [None]*len(train.data)
            pchoices = [None]*len(train.data)
 
            for pno in range(1):
                order = np.random.RandomState(seed=42+pno).permutation(len(train.data))
                for n, ind in enumerate(zip(*(iter(order),)*blocksize)):
                    v = np.array([ np.outer(t, np.append(t, [1])).ravel() for z in ind for t in ( train.data[z], ) ]) 
                    if n == 0 and pno == 0:
                        for i, z in enumerate(ind):
                            if logchoices[z] is None:
                                choice = loggerrand.choice(a=classes, size=1)
                                logchoices[z] = choice[0]
                                pchoices[z] = 1.0 / len(classes)
                    else:
                        predlogp = cls.predict_proba(v)
                        soft = softmax(tau * predlogp, axis=1)

                        for i, z in enumerate(ind):
                            if logchoices[z] is None:
                                choice = loggerrand.choice(a=classes, p=soft[i,:], size=1)
                                logchoices[z] = choice[0]
                                pchoices[z] = soft[i, choice[0]]

                        pred = cls.predict(v)
                        actual = [ train.target[z] for z in ind ]
                        for i, (p, a) in enumerate(zip(pred, actual)):
                            loss += 0 if p == a else 1
                            sincelast += 0 if p == a else 1
                            logpv += soft[i, a]
                            logpvsl += soft[i, a]
                            
                        if (n & (n - 1) == 0): # and n & 0xAAAAAAAA == 0):
                            print('{:<8d}\t{:<8.3f}\t{:<10.3f}\t{:<8.3f}\t{:<10.3f}'.format(
                                        loss.n, loss.mean(), sincelast.mean(), logpv.mean(), logpvsl.mean()),
                                  flush=True)

                            sincelast = OnlineDRO.EasyAcc()
                            logpvsl = OnlineDRO.EasyAcc()

                    x = np.array([ v[i] for i, z in enumerate(ind) if logchoices[z] == train.target[z] ])
                    y = np.array([ logchoices[z] for i, z in enumerate(ind) if logchoices[z] == train.target[z] ])
                    w = np.array([ (lr / len(classes)) * (1 / pchoices[z]) 
                                   for i, z in enumerate(ind) if logchoices[z] == train.target[z] ])
                    if np.any(x):
                        cls.partial_fit(x, y, classes=classes, sample_weight=w)

                print('{:<8d}\t{:<8.3f}\t{:<10.3f}\t{:<8.3f}\t{:<10.3f}'.format(
                             loss.n, loss.mean(), sincelast.mean(), logpv.mean(), logpvsl.mean()),
                       flush=True)                
                sincelast = OnlineDRO.EasyAcc()
                logpvsl = OnlineDRO.EasyAcc()

                preds = cls.predict(np.array([np.outer(d, np.append(d, [1])).ravel() for d in test.data]))
                ascores = []
                for b in range(16):
                    bootie = np.random.RandomState(90210+b).choice(len(test.target), replace=True, size=len(test.target))
                    ascores.append(accuracy_score(y_true=test.target[bootie], y_pred=preds[bootie]))

                print("test accuracy: {}".format(np.quantile(ascores, [0.05, 0.5, 0.95])))
    
OnlineDRO.flass()      


# + [markdown] heading_collapsed=true
# ### Bound Online Learning
#
# Either MLE, lower bound, or upper bound.  Some count decay seems better than no count decay.  Upper bound seems to want less count decay than the MLE or lower bound.  All forms of bound learning have lower regret than IPS learning.
#
#
# TODO: "delayed batch" online learning.

# + code_folding=[1, 14, 17, 34, 52, 59, 66, 74, 82, 90, 141, 150] hidden=true
class OnlineDRO:
    class EasyAcc:
        def __init__(self):
            self.n = 0
            self.sum = 0
            
        def __iadd__(self, other):
            self.n += 1
            self.sum += other
            return self
            
        def mean(self):
            return self.sum / max(self.n, 1)

    class OnlineCressieReadLB:
        from math import inf
        
        def __init__(self, alpha, gamma=1, wmin=0, wmax=inf):
            import numpy as np
            
            self.alpha = alpha
            self.gamma = gamma
            self.n = 0
            self.sumw = 0
            self.sumwsq = 0
            self.sumwr = 0
            self.sumwsqr = 0
            self.sumwsqrsq = 0
            self.wmin = wmin
            self.wmax = wmax
            
            self.duals = None
            self.mleduals = None
            
        def update(self, c, w, r):
            if c > 0:
                assert w + 1e-6 >= self.wmin and w <= self.wmax + 1e-6, 'w = {} < {} < {}'.format(self.wmin, w, self.wmax)
                assert r >= 0 and r <= 1, 'r = {}'.format(r)
                
                decay = self.gamma ** c
                self.n = decay * self.n + c
                self.sumw = decay * self.sumw + c * w
                self.sumwsq = decay * self.sumwsq + c * w**2
                self.sumwr = decay * self.sumwr + c * w * r
                self.sumwsqr = decay * self.sumwsqr + c * (w**2) * r
                self.sumwsqrsq = decay * self.sumwsqrsq + c * (w**2) * (r**2)
                    
                self.duals = None
                self.mleduals = None
                
            return self
        
        def recomputeduals(self):
            from MLE.MLE import CrMinusTwo as CrMinusTwo
            
            self.duals = CrMinusTwo.intervalimpl(self.n, self.sumw, self.sumwsq, 
                                                 self.sumwr, self.sumwsqr, self.sumwsqrsq,
                                                 self.wmin, self.wmax, self.alpha, raiseonerr=True)
            
        def recomputedualsmle(self):
            from MLE.MLE import CrMinusTwo as CrMinusTwo
            
            self.mleduals = CrMinusTwo.estimateimpl(self.n, self.sumw, self.sumwsq, 
                                                    self.sumwr, self.sumwsqr, None, None,
                                                    self.wmin, self.wmax, raiseonerr=True)
        
        def qlb(self, c, w, r):
            if self.duals is None:
                self.recomputeduals()
                
                assert self.duals is not None
                
            return self.duals[1][0]['qfunc'](c, w, r) if self.duals[1][0] is not None else 1
        
        def qub(self, c, w, r):
            if self.duals is None:
                self.recomputeduals()
                
                assert self.duals is not None
                
            return self.duals[1][1]['qfunc'](c, w, r) if self.duals[1][1] is not None else 1
        
        def qmle(self, c, w, r):
            if self.mleduals is None:
                self.recomputedualsmle()
                
                assert self.mleduals is not None
                
            return self.mleduals[1]['qfunc'](c, w, r) if self.mleduals[1] is not None else 1
           
    def flass():
        from scipy.special import softmax
        from sklearn.datasets import fetch_covtype
        from sklearn.decomposition import PCA
        from sklearn.linear_model import SGDClassifier
        from sklearn.metrics import accuracy_score
        from math import ceil, exp
        import numpy as np
        
        cov = fetch_covtype()
        cov.data = PCA(whiten=True).fit_transform(cov.data)
        classes = np.unique(cov.target - 1)
        ndata = len(cov.target)
        order = np.random.RandomState(seed=42).permutation(ndata)
        ntrain = ceil(0.2 * ndata)
        Object = lambda **kwargs: type("Object", (), kwargs)()
        train = Object(data = cov.data[order[:ntrain]], target = cov.target[order[:ntrain]] - 1)
        test = Object(data = cov.data[order[ntrain:]], target = cov.target[order[ntrain:]] - 1)
        
        blocksize = 32
        for lr, tau, what, gamma in ( (x, y, z, g) for x in np.logspace(-2.5, -2, 1) for y in (4, ) 
                                      for z in ('mle', 'ub', 'lb') for g in (1, 0.9999, 0.999, 0.99, ) ):
            print("*** lr = {} tau = {} what = {} gamma = {} ***".format(lr, tau, what, gamma), flush=True)
            print('{:8.8s}\t{:8.8s}\t{:10.10s}\t{:8.8s}\t{:10.10s}\t{:8.8s}\t{:10.10s}'.format(
                'n', 'eff n', 'since last', 'emp loss', 'since last', 'log pv', 'since last')
            )
            
            cls = SGDClassifier(loss='log', shuffle=False)
            loss = OnlineDRO.EasyAcc()
            sincelast = OnlineDRO.EasyAcc()
            logpv = OnlineDRO.EasyAcc()
            logpvsl = OnlineDRO.EasyAcc()
            effn = OnlineDRO.EasyAcc()
            effnsl = OnlineDRO.EasyAcc()
            
            loggerrand = np.random.RandomState(seed=2112)
            logchoices = [None]*len(train.data)
            pchoices = [None]*len(train.data)
                        
            ocrl = OnlineDRO.OnlineCressieReadLB(alpha=0.05, 
                                                 gamma=gamma,
                                                 wmin=0,
                                                 wmax=exp(tau) + len(classes) - 1
                                                )
            qfunc = ocrl.qmle if what == 'mle' else ocrl.qlb if what == 'lb' else ocrl.qub
 
            for pno in range(1):
                order = np.random.RandomState(seed=42+pno).permutation(len(train.data))
                for n, ind in enumerate(zip(*(iter(order),)*blocksize)):
                    v = np.array([ np.outer(t, np.append(t, [1])).ravel() for z in ind for t in ( train.data[z], ) ]) 
                    if n == 0 and pno == 0:
                        for i, z in enumerate(ind):
                            if logchoices[z] is None:
                                choice = loggerrand.choice(a=classes, size=1)
                                logchoices[z] = choice[0]
                                pchoices[z] = 1.0 / len(classes)
                    else:
                        predlogp = cls.predict_proba(v)
                        soft = softmax(tau * predlogp, axis=1)

                        for i, z in enumerate(ind):
                            if logchoices[z] is None:
                                choice = loggerrand.choice(a=classes, p=soft[i,:], size=1)
                                logchoices[z] = choice[0]
                                pchoices[z] = soft[i, choice[0]]

                        pred = cls.predict(v)
                        actual = [ train.target[z] for z in ind ]
                        for i, (p, a) in enumerate(zip(pred, actual)):
                            loss += 0 if p == a else 1
                            sincelast += 0 if p == a else 1
                            logpv += soft[i, a]
                            logpvsl += soft[i, a]
                            
                        if (n & (n - 1) == 0): # and n & 0xAAAAAAAA == 0):
                            print('{:<8d}\t{:<8.3f}\t{:<10.3f}\t{:<8.3f}\t{:<10.3f}\t{:<8.3f}\t{:<10.3f}'.format(
                                         loss.n, effn.mean(), effnsl.mean(), loss.mean(), sincelast.mean(), logpv.mean(), logpvsl.mean()),
                                   flush=True)   

                            sincelast = OnlineDRO.EasyAcc()
                            logpvsl = OnlineDRO.EasyAcc()
                            effnsl = OnlineDRO.EasyAcc()
                            
                    for i, z in enumerate(ind):
                        r = 1 if logchoices[z] == train.target[z] else 0
                        w = 1 / pchoices[z] 
                        ocrl.update(1, w, r) 
                        
                    if n == 0 and pno == 0:
                        sampweight = np.array([ lr for i, z in enumerate(ind) if logchoices[z] == train.target[z] ])
                    else:   
                        sampweight = np.array([ lr * w * ocrl.n 
                                                   * max(0, qfunc(1, w, 1))
                                                for i, z in enumerate(ind) 
                                                if logchoices[z] == train.target[z]
                                                for w in (1 / pchoices[z],)# if logchoices[z] == pred[i] else 0,)
                                              ])
                        
                    effn += sampweight.sum() / (lr * blocksize)
                    effnsl += sampweight.sum() / (lr * blocksize)
                                            
                    x = np.array([ v[i] for i, z in enumerate(ind) if logchoices[z] == train.target[z] ])
                    y = np.array([ logchoices[z] for i, z in enumerate(ind) if logchoices[z] == train.target[z] ])
                    
                    if np.any(x):
                        cls.partial_fit(x, y, classes=classes, sample_weight=sampweight)

                print('{:<8d}\t{:<8.3f}\t{:<10.3f}\t{:<8.3f}\t{:<10.3f}\t{:<8.3f}\t{:<10.3f}'.format(
                             loss.n, effn.mean(), effnsl.mean(), loss.mean(), sincelast.mean(), logpv.mean(), logpvsl.mean()),
                       flush=True)

                sincelast = OnlineDRO.EasyAcc()
                logpvsl = OnlineDRO.EasyAcc()
                effnsl = OnlineDRO.EasyAcc()

                preds = cls.predict(np.array([np.outer(d, np.append(d, [1])).ravel() for d in test.data]))
                ascores = []
                for b in range(16):
                    bootie = np.random.RandomState(90210+b).choice(len(test.target), replace=True, size=len(test.target))
                    ascores.append(accuracy_score(y_true=test.target[bootie], y_pred=preds[bootie]))

                print("test accuracy: {}".format(np.quantile(ascores, [0.05, 0.5, 0.95])))
    
OnlineDRO.flass()      


# + [markdown] heading_collapsed=true
# ### Delayed Batch Bound Online Learning
#
# Process larger batches as sets of smaller batches to emulate time delay in policy updates without changing optimization properties (i.e., SGD batch size).
#
# Everything still seems to work.

# + code_folding=[1, 2, 6, 11, 14, 17, 34, 52, 59, 66, 74, 82, 90, 142] hidden=true
class OnlineDRO:
    class EasyAcc:
        def __init__(self):
            self.n = 0
            self.sum = 0
            
        def __iadd__(self, other):
            self.n += 1
            self.sum += other
            return self
            
        def mean(self):
            return self.sum / max(self.n, 1)

    class OnlineCressieReadLB:
        from math import inf
        
        def __init__(self, alpha, gamma=1, wmin=0, wmax=inf):
            import numpy as np
            
            self.alpha = alpha
            self.gamma = gamma
            self.n = 0
            self.sumw = 0
            self.sumwsq = 0
            self.sumwr = 0
            self.sumwsqr = 0
            self.sumwsqrsq = 0
            self.wmin = wmin
            self.wmax = wmax
            
            self.duals = None
            self.mleduals = None
            
        def update(self, c, w, r):
            if c > 0:
                assert w + 1e-6 >= self.wmin and w <= self.wmax + 1e-6, 'w = {} < {} < {}'.format(self.wmin, w, self.wmax)
                assert r >= 0 and r <= 1, 'r = {}'.format(r)
                
                decay = self.gamma ** c
                self.n = decay * self.n + c
                self.sumw = decay * self.sumw + c * w
                self.sumwsq = decay * self.sumwsq + c * w**2
                self.sumwr = decay * self.sumwr + c * w * r
                self.sumwsqr = decay * self.sumwsqr + c * (w**2) * r
                self.sumwsqrsq = decay * self.sumwsqrsq + c * (w**2) * (r**2)
                    
                self.duals = None
                self.mleduals = None
                
            return self
        
        def recomputeduals(self):
            from MLE.MLE import CrMinusTwo as CrMinusTwo
            
            self.duals = CrMinusTwo.intervalimpl(self.n, self.sumw, self.sumwsq, 
                                                 self.sumwr, self.sumwsqr, self.sumwsqrsq,
                                                 self.wmin, self.wmax, self.alpha, raiseonerr=True)
            
        def recomputedualsmle(self):
            from MLE.MLE import CrMinusTwo as CrMinusTwo
            
            self.mleduals = CrMinusTwo.estimateimpl(self.n, self.sumw, self.sumwsq, 
                                                    self.sumwr, self.sumwsqr, None, None,
                                                    self.wmin, self.wmax, raiseonerr=True)
        
        def qlb(self, c, w, r):
            if self.duals is None:
                self.recomputeduals()
                
                assert self.duals is not None
                
            return self.duals[1][0]['qfunc'](c, w, r) if self.duals[1][0] is not None else 1
        
        def qub(self, c, w, r):
            if self.duals is None:
                self.recomputeduals()
                
                assert self.duals is not None
                
            return self.duals[1][1]['qfunc'](c, w, r) if self.duals[1][1] is not None else 1
        
        def qmle(self, c, w, r):
            if self.mleduals is None:
                self.recomputedualsmle()
                
                assert self.mleduals is not None
                
            return self.mleduals[1]['qfunc'](c, w, r) if self.mleduals[1] is not None else 1
           
    def flass():
        from scipy.special import softmax
        from sklearn.datasets import fetch_covtype
        from sklearn.decomposition import PCA
        from sklearn.linear_model import SGDClassifier
        from sklearn.metrics import accuracy_score
        from math import ceil, exp
        import numpy as np
        
        cov = fetch_covtype()
        cov.data = PCA(whiten=True).fit_transform(cov.data)
        classes = np.unique(cov.target - 1)
        ndata = len(cov.target)
        order = np.random.RandomState(seed=42).permutation(ndata)
        ntrain = ceil(0.2 * ndata)
        Object = lambda **kwargs: type("Object", (), kwargs)()
        train = Object(data = cov.data[order[:ntrain]], target = cov.target[order[:ntrain]] - 1)
        test = Object(data = cov.data[order[ntrain:]], target = cov.target[order[ntrain:]] - 1)
        
        subblocksize = 32
        delay = 8
        blocksize = delay * subblocksize
        for lr, tau, what, gamma in ( (x, y, z, g) for x in np.logspace(-2.5, -2, 1) for y in (4, ) 
                                      for z in ('mle', 'ub', 'lb') for g in (1, 0.9999, 0.999, 0.99,) ):
            print("*** lr = {} tau = {} what = {} gamma = {} ***".format(lr, tau, what, gamma), flush=True)
            print('{:8.8s}\t{:8.8s}\t{:10.10s}\t{:8.8s}\t{:10.10s}\t{:8.8s}\t{:10.10s}'.format(
                'n', 'eff n', 'since last', 'emp loss', 'since last', 'log pv', 'since last')
            )
            
            cls = SGDClassifier(loss='log', shuffle=False)
            loss = OnlineDRO.EasyAcc()
            sincelast = OnlineDRO.EasyAcc()
            logpv = OnlineDRO.EasyAcc()
            logpvsl = OnlineDRO.EasyAcc()
            effn = OnlineDRO.EasyAcc()
            effnsl = OnlineDRO.EasyAcc()
            
            loggerrand = np.random.RandomState(seed=2112)
            logchoices = [None]*len(train.data)
            pchoices = [None]*len(train.data)
                        
            ocrl = OnlineDRO.OnlineCressieReadLB(alpha=0.05, 
                                                 gamma=gamma,
                                                 wmin=0,
                                                 wmax=exp(tau) + len(classes) - 1
                                                )
            qfunc = ocrl.qmle if what == 'mle' else ocrl.qlb if what == 'lb' else ocrl.qub
 
            for pno in range(1):
                order = np.random.RandomState(seed=42+pno).permutation(len(train.data))
                for n, ind in enumerate(zip(*(iter(order),)*blocksize)):
                    v = np.array([ np.outer(t, np.append(t, [1])).ravel() for z in ind for t in ( train.data[z], ) ]) 
                    if n == 0 and pno == 0:
                        for i, z in enumerate(ind):
                            if logchoices[z] is None:
                                choice = loggerrand.choice(a=classes, size=1)
                                logchoices[z] = choice[0]
                                pchoices[z] = 1.0 / len(classes)
                    else:
                        predlogp = cls.predict_proba(v)
                        soft = softmax(tau * predlogp, axis=1)

                        for i, z in enumerate(ind):
                            if logchoices[z] is None:
                                choice = loggerrand.choice(a=classes, p=soft[i,:], size=1)
                                logchoices[z] = choice[0]
                                pchoices[z] = soft[i, choice[0]]

                        pred = cls.predict(v)
                        actual = [ train.target[z] for z in ind ]
                        for i, (p, a) in enumerate(zip(pred, actual)):
                            loss += 0 if p == a else 1
                            sincelast += 0 if p == a else 1
                            logpv += soft[i, a]
                            logpvsl += soft[i, a]
                            
                        if (n & (n - 1) == 0): # and n & 0xAAAAAAAA == 0):
                            print('{:<8d}\t{:<8.3f}\t{:<10.3f}\t{:<8.3f}\t{:<10.3f}\t{:<8.3f}\t{:<10.3f}'.format(
                                         loss.n, effn.mean(), effnsl.mean(), loss.mean(), sincelast.mean(), logpv.mean(), logpvsl.mean()),
                                   flush=True)   

                            sincelast = OnlineDRO.EasyAcc()
                            logpvsl = OnlineDRO.EasyAcc()
                            effnsl = OnlineDRO.EasyAcc()
                            
                    for i, z in enumerate(ind):
                        r = 1 if logchoices[z] == train.target[z] else 0
                        w = 1 / pchoices[z] 
                        ocrl.update(1, w, r) 
                                                   
                    for d in range(delay):
                        x = np.array([ v[i] 
                                       for i, z in enumerate(ind)
                                       if (d-1)*subblocksize <= i and i < d*subblocksize
                                       if logchoices[z] == train.target[z] ])
                        y = np.array([ logchoices[z] 
                                       for i, z in enumerate(ind) 
                                       if (d-1)*subblocksize <= i and i < d*subblocksize
                                       if logchoices[z] == train.target[z] ])
                        
                        if n == 0 and pno == 0:
                            sampweight = np.array([ lr 
                                                    for i, z in enumerate(ind) 
                                                    if (d-1)*subblocksize <= i and i < d*subblocksize
                                                    if logchoices[z] == train.target[z] ])
                        else:
                            sampweight = np.array([ lr * w * ocrl.n 
                                                       * max(0, qfunc(1, w, 1))
                                                    for i, z in enumerate(ind) 
                                                    if (d-1)*subblocksize <= i and i < d*subblocksize
                                                    if logchoices[z] == train.target[z]
                                                    for w in (1 / pchoices[z],)
                                                  ])
                        
                        effn += sampweight.sum() / (lr * subblocksize)
                        effnsl += sampweight.sum() / (lr * subblocksize)
                                             
                        if np.any(x):
                            cls.partial_fit(x, y, classes=classes, sample_weight=sampweight)

                print('{:<8d}\t{:<8.3f}\t{:<10.3f}\t{:<8.3f}\t{:<10.3f}\t{:<8.3f}\t{:<10.3f}'.format(
                             loss.n, effn.mean(), effnsl.mean(), loss.mean(), sincelast.mean(), logpv.mean(), logpvsl.mean()),
                       flush=True)
                
#                 from pprint import pformat
#                 print(pformat(ocrl.__dict__))

                sincelast = OnlineDRO.EasyAcc()
                logpvsl = OnlineDRO.EasyAcc()
                effnsl = OnlineDRO.EasyAcc()

                preds = cls.predict(np.array([np.outer(d, np.append(d, [1])).ravel() for d in test.data]))
                ascores = []
                for b in range(16):
                    bootie = np.random.RandomState(90210+b).choice(len(test.target), replace=True, size=len(test.target))
                    ascores.append(accuracy_score(y_true=test.target[bootie], y_pred=preds[bootie]))

                print("test accuracy: {}".format(np.quantile(ascores, [0.05, 0.5, 0.95])))
    
OnlineDRO.flass()      


# -

# ### Auto-Temperature + Delayed Batch Bound Online Learning
#
# Best regret achieved with lower bound optimization.

# + code_folding=[1, 90, 118, 197, 203, 217, 257]
class OnlineDRO:
    class EasyAcc:
        def __init__(self):
            self.n = 0
            self.sum = 0
            
        def __iadd__(self, other):
            self.n += 1
            self.sum += other
            return self
            
        def mean(self):
            return self.sum / max(self.n, 1)

    class OnlineCressieReadLB:
        from math import inf
        
        def __init__(self, alpha, gamma=1, wmin=0, wmax=inf):
            import numpy as np
            
            self.alpha = alpha
            self.gamma = gamma
            self.n = 0
            self.sumw = 0
            self.sumwsq = 0
            self.sumwr = 0
            self.sumwsqr = 0
            self.sumwsqrsq = 0
            self.wmin = wmin
            self.wmax = wmax
            
            self.duals = None
            self.mleduals = None
            
        def update(self, c, w, r):
            if c > 0:
                assert w + 1e-6 >= self.wmin and w <= self.wmax + 1e-6, 'w = {} < {} < {}'.format(self.wmin, w, self.wmax)
                assert r >= 0 and r <= 1, 'r = {}'.format(r)
                
                decay = self.gamma ** c
                self.n = decay * self.n + c
                self.sumw = decay * self.sumw + c * w
                self.sumwsq = decay * self.sumwsq + c * w**2
                self.sumwr = decay * self.sumwr + c * w * r
                self.sumwsqr = decay * self.sumwsqr + c * (w**2) * r
                self.sumwsqrsq = decay * self.sumwsqrsq + c * (w**2) * (r**2)
                    
                self.duals = None
                self.mleduals = None
                
            return self
        
        def recomputeduals(self):
            from MLE.MLE import CrMinusTwo as CrMinusTwo
            
            self.duals = CrMinusTwo.intervalimpl(self.n, self.sumw, self.sumwsq, 
                                                 self.sumwr, self.sumwsqr, self.sumwsqrsq,
                                                 self.wmin, self.wmax, self.alpha, raiseonerr=True)
            
        def recomputedualsmle(self):
            from MLE.MLE import CrMinusTwo as CrMinusTwo
            
            self.mleduals = CrMinusTwo.estimateimpl(self.n, self.sumw, self.sumwsq, 
                                                    self.sumwr, self.sumwsqr, None, None,
                                                    self.wmin, self.wmax, raiseonerr=True)
        
        def qlb(self, c, w, r):
            if self.duals is None:
                self.recomputeduals()
                
                assert self.duals is not None
                
            return self.duals[1][0]['qfunc'](c, w, r) if self.duals[1][0] is not None else 1
        
        def qub(self, c, w, r):
            if self.duals is None:
                self.recomputeduals()
                
                assert self.duals is not None
                
            return self.duals[1][1]['qfunc'](c, w, r) if self.duals[1][1] is not None else 1
        
        def qmle(self, c, w, r):
            if self.mleduals is None:
                self.recomputedualsmle()
                
                assert self.mleduals is not None
                
            return self.mleduals[1]['qfunc'](c, w, r) if self.mleduals[1] is not None else 1
           
    def autotune(pre, target, taumax):
        from scipy.optimize import root_scalar
        
        def f(tau):
            from scipy.special import softmax
            import numpy as np

            soft = softmax(tau * pre, axis=1)
            minsoft = np.min(soft, axis=1)
            
            return np.mean(minsoft) - target
            
        fmax = f(taumax)
        
        if fmax >= 0:
            return taumax
        
        taumin = 0
        fmin = f(taumin)
        
        assert fmin > 0, { 'fmin': fmin, 'target': target }
        
        root = root_scalar(f, bracket=(taumin, taumax))
        
        assert root.converged, root
                
        return root.root
            
    def flass():
        from scipy.special import softmax
        from sklearn.datasets import fetch_covtype
        from sklearn.decomposition import PCA
        from sklearn.linear_model import SGDClassifier
        from sklearn.metrics import accuracy_score
        from math import ceil, exp
        import numpy as np
        
        cov = fetch_covtype()
        cov.data = PCA(whiten=True).fit_transform(cov.data)
        classes = np.unique(cov.target - 1)
        ndata = len(cov.target)
        order = np.random.RandomState(seed=42).permutation(ndata)
        ntrain = ceil(0.2 * ndata)
        Object = lambda **kwargs: type("Object", (), kwargs)()
        train = Object(data = cov.data[order[:ntrain]], target = cov.target[order[:ntrain]] - 1)
        test = Object(data = cov.data[order[ntrain:]], target = cov.target[order[ntrain:]] - 1)
        
        subblocksize = 32
        delay = 8
        blocksize = delay * subblocksize
        for lr, taumax, target, what, gamma in ( (x, 8, y, z, g) for x in np.logspace(-2.5, -2, 1) 
                                         for z in ('mle', 'ub', 'lb') for g in (0.9999, 0.999,) for y in (0.025, 0.05, 0.1, 0.2, ) ):
            print("*** lr = {} taumax = {} target = {} what = {} gamma = {} ***".format(lr, taumax, target, what, gamma), flush=True)
            print('{:8.8s}\t{:8.8s} [{:10.10s}]\t{:8.8s} [{:10.10s}]\t{:8.8s} [{:10.10s}]\t{:8.8s} [{:10.10s}]'.format(
                'n', 'eff n', 'since last', 'av tau', 'since last', 'emp loss', 'since last', 'log pv', 'since last')
            )
            
            try:
                target /= len(classes)

                cls = SGDClassifier(loss='log', shuffle=False)
                loss = OnlineDRO.EasyAcc()
                sincelast = OnlineDRO.EasyAcc()
                logpv = OnlineDRO.EasyAcc()
                logpvsl = OnlineDRO.EasyAcc()
                effn = OnlineDRO.EasyAcc()
                effnsl = OnlineDRO.EasyAcc()
                avtau = OnlineDRO.EasyAcc()
                avtausl = OnlineDRO.EasyAcc()

                loggerrand = np.random.RandomState(seed=2112)
                logchoices = [None]*len(train.data)
                pchoices = [None]*len(train.data)

                ocrl = OnlineDRO.OnlineCressieReadLB(alpha=0.05, 
                                                     gamma=gamma,
                                                     wmin=0,
#                                                      wmax=exp(taumax) + len(classes) - 1
                                                    )
                qfunc = ocrl.qmle if what == 'mle' else ocrl.qlb if what == 'lb' else ocrl.qub

                for pno in range(1):
                    order = np.random.RandomState(seed=42+pno).permutation(len(train.data))
                    for n, ind in enumerate(zip(*(iter(order),)*blocksize)):
                        v = np.array([ np.outer(t, np.append(t, [1])).ravel() for z in ind for t in ( train.data[z], ) ]) 
                        if n == 0 and pno == 0:
                            pred = np.zeros(blocksize)
                            for i, z in enumerate(ind):
                                if logchoices[z] is None:
                                    choice = loggerrand.choice(a=classes, size=1)
                                    logchoices[z] = choice[0]
                                    pchoices[z] = 1.0 / len(classes)
                        else:
                            predlogp = cls.predict_proba(v)
                            tau = OnlineDRO.autotune(predlogp, target, taumax)
                            avtau += tau
                            avtausl += tau
                            soft = softmax(tau * predlogp, axis=1)

                            for i, z in enumerate(ind):
                                if logchoices[z] is None:
                                    choice = loggerrand.choice(a=classes, p=soft[i,:], size=1)
                                    logchoices[z] = choice[0]
                                    pchoices[z] = soft[i, choice[0]]

                            pred = cls.predict(v)
                            actual = [ train.target[z] for z in ind ]
                            for i, (p, a) in enumerate(zip(pred, actual)):
                                loss += 0 if p == a else 1
                                sincelast += 0 if p == a else 1
                                logpv += soft[i, a]
                                logpvsl += soft[i, a]

                            if (n & (n - 1) == 0): # and n & 0xAAAAAAAA == 0):
                                print('{:<8d}\t{:<8.3f} [{:<10.3f}]\t{:<8.3f} [{:<10.3f}]\t{:<8.3f} [{:<10.3f}]\t{:<8.3f} [{:<10.3f}]'.format(
                                             loss.n, effn.mean(), effnsl.mean(), avtau.mean(), avtausl.mean(), loss.mean(), sincelast.mean(), logpv.mean(), logpvsl.mean()),
                                       flush=True)   
                                sincelast = OnlineDRO.EasyAcc()
                                logpvsl = OnlineDRO.EasyAcc()
                                effnsl = OnlineDRO.EasyAcc()
                                avtausl = OnlineDRO.EasyAcc()

                        for i, z in enumerate(ind):
                            r = 1 if logchoices[z] == train.target[z] else 0
                            w = 1 / pchoices[z] if pred[i] == logchoices[z] else 0
                            ocrl.update(1, w, r) 

                        for d in range(delay):
                            x = np.array([ v[i] 
                                           for i, z in enumerate(ind)
                                           if (d-1)*subblocksize <= i and i < d*subblocksize
                                           if logchoices[z] == train.target[z] ])
                            y = np.array([ logchoices[z] 
                                           for i, z in enumerate(ind) 
                                           if (d-1)*subblocksize <= i and i < d*subblocksize
                                           if logchoices[z] == train.target[z] ])

                            if n == 0 and pno == 0:
                                sampweight = np.array([ lr 
                                                        for i, z in enumerate(ind) 
                                                        if (d-1)*subblocksize <= i and i < d*subblocksize
                                                        if logchoices[z] == train.target[z] ])
                            else:
                                sampweight = np.array([ lr * w * ocrl.n 
                                                           * max(0, qfunc(1, w, 1))
                                                        for i, z in enumerate(ind) 
                                                        if (d-1)*subblocksize <= i and i < d*subblocksize
                                                        if logchoices[z] == train.target[z]
                                                        for w in (1 / pchoices[z],)
                                                      ])

                            effn += sampweight.sum() / (lr * subblocksize)
                            effnsl += sampweight.sum() / (lr * subblocksize)

                            if np.any(x):
                                cls.partial_fit(x, y, classes=classes, sample_weight=sampweight)

                    print('{:<8d}\t{:<8.3f} [{:<10.3f}]\t{:<8.3f} [{:<10.3f}]\t{:<8.3f} [{:<10.3f}]\t{:<8.3f} [{:<10.3f}]'.format(
                                     loss.n, effn.mean(), effnsl.mean(), avtau.mean(), avtausl.mean(), loss.mean(), sincelast.mean(), logpv.mean(), logpvsl.mean()),
                               flush=True)   
                    sincelast = OnlineDRO.EasyAcc()
                    logpvsl = OnlineDRO.EasyAcc()
                    effnsl = OnlineDRO.EasyAcc()
                    avtausl = OnlineDRO.EasyAcc()

                    preds = cls.predict(np.array([np.outer(d, np.append(d, [1])).ravel() for d in test.data]))
                    ascores = []
                    for b in range(16):
                        bootie = np.random.RandomState(90210+b).choice(len(test.target), replace=True, size=len(test.target))
                        ascores.append(accuracy_score(y_true=test.target[bootie], y_pred=preds[bootie]))

                    print("test accuracy: {}".format(np.quantile(ascores, [0.05, 0.5, 0.95])))
            except KeyboardInterrupt:
                raise
            except:
                pass
    
OnlineDRO.flass()      


# + [markdown] heading_collapsed=true
# # 20 Newsgroups
#
# Really hard ... 20 actions and only 22K examples.

# + [markdown] hidden=true
# ### Full Information Online Learning

# + code_folding=[1, 14] hidden=true
class OnlineDRO:
    class EasyAcc:
        def __init__(self):
            self.n = 0
            self.sum = 0
            
        def __iadd__(self, other):
            self.n += 1
            self.sum += other
            return self
            
        def mean(self):
            return self.sum / max(self.n, 1)
 
    def flass():
        from scipy.sparse import vstack
        from sklearn.datasets import fetch_20newsgroups
        from sklearn.feature_extraction.text import HashingVectorizer
        from sklearn.linear_model import SGDClassifier
        from sklearn.metrics import accuracy_score
        import numpy as np
        
        train = fetch_20newsgroups(subset='train', remove=('headers', 'footers', 'quotes'))
        classes = np.unique(train.target)
        test = fetch_20newsgroups(subset='test', remove=('headers', 'footers', 'quotes'))
        
        for ngram, norm, lr, nfeat in ( (2, 'l2', 1, 20), ):
            print("**** ngram = {} norm = {} lr = {} nfeat = {} ****".format(ngram, norm, lr, nfeat))
            print('{:8.8s}\t{:8.8s}\t{:10.10s}'.format('n', 'emp loss', 'since last'))

            vectorizer = HashingVectorizer(n_features = 1 << nfeat, norm=norm, ngram_range=(1, ngram), alternate_sign=True)
            docs = vectorizer.transform(train.data)
            testdocs = vectorizer.transform(test.data)
            
            classweights = { k: lr for k in classes }
            cls = SGDClassifier(loss='log', class_weight=classweights, shuffle=False)
            
            loss = OnlineDRO.EasyAcc()
            sincelast = OnlineDRO.EasyAcc()
            blocksize = 32

            for pno in range(1):
                order = np.random.RandomState(seed=42+pno).permutation(len(train.data))
                for n, ind in enumerate(zip(*(iter(order),)*blocksize)):
                    v = vstack([ docs[z] for z in ind ])
                    actual = [ train.target[z] for z in ind ]
                    if n > 0:
                        pred = cls.predict(v)
                        for p, a in zip(pred, actual):
                            loss += 0 if p == a else 1
                            sincelast += 0 if p == a else 1
                        if (n & (n - 1) == 0): # and n & 0xAAAAAAAA == 0):
                            print('{:<8d}\t{:<8.3f}\t{:<10.3f}'.format(loss.n, loss.mean(), sincelast.mean()), flush=True)
                            sincelast = OnlineDRO.EasyAcc()

                    cls.partial_fit(v, actual, classes=classes)

                print('{:<8d}\t{:<8.3f}\t{:<10.3f}'.format(loss.n, loss.mean(), sincelast.mean()), flush=True)
                sincelast = OnlineDRO.EasyAcc()

                preds = cls.predict(testdocs)
                ascores = []
                for b in range(16):
                    bootie = np.random.RandomState(90210+b).choice(len(test.target), replace=True, size=len(test.target))
                    ascores.append(accuracy_score(y_true=test.target[bootie], y_pred=preds[bootie]))

                print("test accuracy: {}".format(np.quantile(ascores, [0.05, 0.5, 0.95])))
    
OnlineDRO.flass()


# + [markdown] hidden=true
# ### Partial Information Online Learning, Softmax Logging Policy
#
# Uniform $(\tau = 0)$ and softmax $(\tau = 4)$ are pretty similar for off-policy learning but uniform has larger regret.

# + code_folding=[1, 14] hidden=true
class OnlineDRO:
    class EasyAcc:
        def __init__(self):
            self.n = 0
            self.sum = 0
            
        def __iadd__(self, other):
            self.n += 1
            self.sum += other
            return self
            
        def mean(self):
            return self.sum / max(self.n, 1)
 
    def flass():
        from scipy.sparse import vstack
        from scipy.special import softmax
        from sklearn.datasets import fetch_20newsgroups
        from sklearn.feature_extraction.text import HashingVectorizer
        from sklearn.linear_model import SGDClassifier
        from sklearn.metrics import accuracy_score
        import numpy as np
        
        train = fetch_20newsgroups(subset='train', remove=('headers', 'footers', 'quotes'))
        classes = np.unique(train.target)
        test = fetch_20newsgroups(subset='test', remove=('headers', 'footers', 'quotes'))
        
        blocksize = 32
        for ngram, norm, lr, nfeat, tau in ( (2, 'l2', z, 20, x) for x in (0, 4,) for z in (0.5, ) ):
            print("**** ngram = {} norm = {} lr = {} nfeat = {} tau = {} ****".format(ngram, norm, lr, nfeat, tau))
            print('{:8.8s}\t{:8.8s}\t{:10.10s}\t{:8.8s}\t{:10.10s}'.format(
                'n', 'emp loss', 'since last', 'log pv', 'since last')
            )
            
            vectorizer = HashingVectorizer(n_features = 1 << nfeat, norm=norm, ngram_range=(1, ngram), alternate_sign=True)
            docs = vectorizer.transform(train.data)
            testdocs = vectorizer.transform(test.data)
            
            classweights = { k: lr for k in classes }
            cls = SGDClassifier(loss='log', class_weight=classweights, shuffle=False)
            
            loss = OnlineDRO.EasyAcc()
            sincelast = OnlineDRO.EasyAcc()
            logpv = OnlineDRO.EasyAcc()
            logpvsl = OnlineDRO.EasyAcc()
            
            loggerrand = np.random.RandomState(seed=2112)
            logchoices = [None]*len(train.data)
            pchoices = [None]*len(train.data)
 
            for pno in range(1):
                order = np.random.RandomState(seed=42+pno).permutation(len(train.data))
                for n, ind in enumerate(zip(*(iter(order),)*blocksize)):
                    v = vstack([ docs[z] for z in ind ])
                    if n == 0 and pno == 0:
                        for i, z in enumerate(ind):
                            if logchoices[z] is None:
                                choice = loggerrand.choice(a=classes, size=1)
                                logchoices[z] = choice[0]
                                pchoices[z] = 1.0 / len(classes)
                    else:
                        predlogp = cls.predict_proba(v)
                        soft = softmax(tau * predlogp, axis=1)

                        for i, z in enumerate(ind):
                            if logchoices[z] is None:
                                choice = loggerrand.choice(a=classes, p=soft[i,:], size=1)
                                logchoices[z] = choice[0]
                                pchoices[z] = soft[i, choice[0]]

                        pred = cls.predict(v)
                        actual = [ train.target[z] for z in ind ]
                        for i, (p, a) in enumerate(zip(pred, actual)):
                            loss += 0 if p == a else 1
                            sincelast += 0 if p == a else 1
                            logpv += soft[i, a]
                            logpvsl += soft[i, a]
                            
                        if (n & (n - 1) == 0): # and n & 0xAAAAAAAA == 0):
                            print('{:<8d}\t{:<8.3f}\t{:<10.3f}\t{:<8.3f}\t{:<10.3f}'.format(
                                        loss.n, loss.mean(), sincelast.mean(), logpv.mean(), logpvsl.mean()),
                                  flush=True)

                            sincelast = OnlineDRO.EasyAcc()
                            logpvsl = OnlineDRO.EasyAcc()

#                     y = np.array([ logchoices[z] for i, z in enumerate(ind) ])
#                     w = np.array([ (lr / len(classes)) * (1 / pchoices[z]) 
#                                     if logchoices[z] == train.target[z]
#                                     else -(lr / len(classes)**2) * (1/pchoices[z])
#                                    for i, z in enumerate(ind) ])
                    
                    y = np.array([ logchoices[z] for i, z in enumerate(ind) if logchoices[z] == train.target[z] ])
                    w = np.array([ (lr / len(classes)) * (1 / pchoices[z]) 
                                   for i, z in enumerate(ind) if logchoices[z] == train.target[z] ])
                    if np.any(y):
                        x = vstack([ v[i] for i, z in enumerate(ind) if logchoices[z] == train.target[z] ])
#                         x = vstack([ v[i] for i, z in enumerate(ind) ])
                        cls.partial_fit(x, y, classes=classes, sample_weight=w)

                print('{:<8d}\t{:<8.3f}\t{:<10.3f}\t{:<8.3f}\t{:<10.3f}'.format(
                             loss.n, loss.mean(), sincelast.mean(), logpv.mean(), logpvsl.mean()),
                       flush=True)                
                sincelast = OnlineDRO.EasyAcc()
                logpvsl = OnlineDRO.EasyAcc()

                preds = cls.predict(testdocs)
                ascores = []
                for b in range(16):
                    bootie = np.random.RandomState(90210+b).choice(len(test.target), replace=True, size=len(test.target))
                    ascores.append(accuracy_score(y_true=test.target[bootie], y_pred=preds[bootie]))

                print("test accuracy: {}".format(np.quantile(ascores, [0.05, 0.5, 0.95])))
    
OnlineDRO.flass()      


# + [markdown] hidden=true
# ### Auto-Temperature and Bound Optimization
#
# Nothing working great, but regret is better than other approaches. 

# + code_folding=[0, 1, 14, 90] hidden=true
class OnlineDRO:
    class EasyAcc:
        def __init__(self):
            self.n = 0
            self.sum = 0
            
        def __iadd__(self, other):
            self.n += 1
            self.sum += other
            return self
            
        def mean(self):
            return self.sum / max(self.n, 1)

    class OnlineCressieReadLB:
        from math import inf
        
        def __init__(self, alpha, gamma=1, wmin=0, wmax=inf):
            import numpy as np
            
            self.alpha = alpha
            self.gamma = gamma
            self.n = 0
            self.sumw = 0
            self.sumwsq = 0
            self.sumwr = 0
            self.sumwsqr = 0
            self.sumwsqrsq = 0
            self.wmin = wmin
            self.wmax = wmax
            
            self.duals = None
            self.mleduals = None
            
        def update(self, c, w, r):
            if c > 0:
                assert w + 1e-6 >= self.wmin and w <= self.wmax + 1e-6, 'w = {} < {} < {}'.format(self.wmin, w, self.wmax)
                assert r >= 0 and r <= 1, 'r = {}'.format(r)
                
                decay = self.gamma ** c
                self.n = decay * self.n + c
                self.sumw = decay * self.sumw + c * w
                self.sumwsq = decay * self.sumwsq + c * w**2
                self.sumwr = decay * self.sumwr + c * w * r
                self.sumwsqr = decay * self.sumwsqr + c * (w**2) * r
                self.sumwsqrsq = decay * self.sumwsqrsq + c * (w**2) * (r**2)
                    
                self.duals = None
                self.mleduals = None
                
            return self
        
        def recomputeduals(self):
            from MLE.MLE import CrMinusTwo as CrMinusTwo
            
            self.duals = CrMinusTwo.intervalimpl(self.n, self.sumw, self.sumwsq, 
                                                 self.sumwr, self.sumwsqr, self.sumwsqrsq,
                                                 self.wmin, self.wmax, self.alpha, raiseonerr=True)
            
        def recomputedualsmle(self):
            from MLE.MLE import CrMinusTwo as CrMinusTwo
            
            self.mleduals = CrMinusTwo.estimateimpl(self.n, self.sumw, self.sumwsq, 
                                                    self.sumwr, self.sumwsqr, None, None,
                                                    self.wmin, self.wmax, raiseonerr=True)
        
        def qlb(self, c, w, r):
            if self.duals is None:
                self.recomputeduals()
                
                assert self.duals is not None
                
            return self.duals[1][0]['qfunc'](c, w, r) if self.duals[1][0] is not None else 1
        
        def qub(self, c, w, r):
            if self.duals is None:
                self.recomputeduals()
                
                assert self.duals is not None
                
            return self.duals[1][1]['qfunc'](c, w, r) if self.duals[1][1] is not None else 1
        
        def qmle(self, c, w, r):
            if self.mleduals is None:
                self.recomputedualsmle()
                
                assert self.mleduals is not None
                
            return self.mleduals[1]['qfunc'](c, w, r) if self.mleduals[1] is not None else 1
           
    def autotune(pre, target, taumax):
        from scipy.optimize import root_scalar
        
        def f(tau):
            from scipy.special import softmax
            import numpy as np

            soft = softmax(tau * pre, axis=1)
            minsoft = np.min(soft, axis=1)
            
            return np.mean(minsoft) - target
            
        fmax = f(taumax)
        
        if fmax >= 0:
            return taumax
        
        taumin = 0
        fmin = f(taumin)
        
        assert fmin > 0, { 'fmin': fmin, 'target': target }
        
        root = root_scalar(f, bracket=(taumin, taumax))
        
        assert root.converged, root
                
        return root.root
            
    def flass():
        from scipy.sparse import vstack
        from scipy.special import softmax
        from sklearn.datasets import fetch_20newsgroups
        from sklearn.feature_extraction.text import HashingVectorizer
        from sklearn.linear_model import SGDClassifier
        from sklearn.metrics import accuracy_score
        from math import exp
        import numpy as np
        
        train = fetch_20newsgroups(subset='train', remove=('headers', 'footers', 'quotes'))
        classes = np.unique(train.target)
        test = fetch_20newsgroups(subset='test', remove=('headers', 'footers', 'quotes'))
                
        subblocksize = 32
        delay = 8
        blocksize = delay * subblocksize
        for lr, taumax, target, what, gamma in ( (x, 8, y, z, g) for x in (0.5,)
                                                                 for z in ('mle', 'ub', 'lb') 
                                                                 for g in (0.9999, 0.999,) 
                                                                 for y in (0.025, 0.05, 0.1, 0.2, ) ):
            
            ngram = 2
            norm = 'l2'
            nfeat = 20
            
            print("*** lr = {} taumax = {} target = {} what = {} gamma = {} ***".format(lr, taumax, target, what, gamma), flush=True)
            print('{:8.8s}\t{:8.8s} [{:10.10s}]\t{:8.8s} [{:10.10s}]\t{:8.8s} [{:10.10s}]\t{:8.8s} [{:10.10s}]'.format(
                'n', 'eff n', 'since last', 'av tau', 'since last', 'emp loss', 'since last', 'log pv', 'since last')
            )
            
            try:
                target /= len(classes)

                vectorizer = HashingVectorizer(n_features = 1 << nfeat, norm=norm, ngram_range=(1, ngram), alternate_sign=True)
                docs = vectorizer.transform(train.data)
                testdocs = vectorizer.transform(test.data)
            
                classweights = { k: lr for k in classes }
                cls = SGDClassifier(loss='log', class_weight=classweights, shuffle=False)
                
                loss = OnlineDRO.EasyAcc()
                sincelast = OnlineDRO.EasyAcc()
                logpv = OnlineDRO.EasyAcc()
                logpvsl = OnlineDRO.EasyAcc()
                effn = OnlineDRO.EasyAcc()
                effnsl = OnlineDRO.EasyAcc()
                avtau = OnlineDRO.EasyAcc()
                avtausl = OnlineDRO.EasyAcc()

                loggerrand = np.random.RandomState(seed=2112)
                logchoices = [None]*len(train.data)
                pchoices = [None]*len(train.data)

                ocrl = OnlineDRO.OnlineCressieReadLB(alpha=0.05, 
                                                     gamma=gamma,
                                                     wmin=0,
                                                     wmax=exp(taumax) + len(classes) - 1
                                                    )
                qfunc = ocrl.qmle if what == 'mle' else ocrl.qlb if what == 'lb' else ocrl.qub

                for pno in range(1):
                    order = np.random.RandomState(seed=42+pno).permutation(len(train.data))
                    for n, ind in enumerate(zip(*(iter(order),)*blocksize)):
                        v = vstack([ docs[z] for z in ind ])
                        if n == 0 and pno == 0:
                            pred = np.zeros(blocksize)
                            for i, z in enumerate(ind):
                                if logchoices[z] is None:
                                    choice = loggerrand.choice(a=classes, size=1)
                                    logchoices[z] = choice[0]
                                    pchoices[z] = 1.0 / len(classes)
                        else:
                            predlogp = cls.predict_proba(v)
                            tau = OnlineDRO.autotune(predlogp, target, taumax)
                            avtau += tau
                            avtausl += tau
                            soft = softmax(tau * predlogp, axis=1)

                            for i, z in enumerate(ind):
                                if logchoices[z] is None:
                                    choice = loggerrand.choice(a=classes, p=soft[i,:], size=1)
                                    logchoices[z] = choice[0]
                                    pchoices[z] = soft[i, choice[0]]

                            pred = cls.predict(v)
                            actual = [ train.target[z] for z in ind ]
                            for i, (p, a) in enumerate(zip(pred, actual)):
                                loss += 0 if p == a else 1
                                sincelast += 0 if p == a else 1
                                logpv += soft[i, a]
                                logpvsl += soft[i, a]

                            if (n & (n - 1) == 0): # and n & 0xAAAAAAAA == 0):
                                print('{:<8d}\t{:<8.3f} [{:<10.3f}]\t{:<8.3f} [{:<10.3f}]\t{:<8.3f} [{:<10.3f}]\t{:<8.3f} [{:<10.3f}]'.format(
                                             loss.n, effn.mean(), effnsl.mean(), avtau.mean(), avtausl.mean(), loss.mean(), sincelast.mean(), logpv.mean(), logpvsl.mean()),
                                       flush=True)   
                                sincelast = OnlineDRO.EasyAcc()
                                logpvsl = OnlineDRO.EasyAcc()
                                effnsl = OnlineDRO.EasyAcc()
                                avtausl = OnlineDRO.EasyAcc()

                        for i, z in enumerate(ind):
                            r = 1 if logchoices[z] == train.target[z] else 0
                            w = 1 / pchoices[z] if pred[i] == logchoices[z] else 0
                            ocrl.update(1, w, r) 

                        for d in range(delay):
                            y = np.array([ logchoices[z] 
                                           for i, z in enumerate(ind) 
                                           if (d-1)*subblocksize <= i and i < d*subblocksize
                                           if logchoices[z] == train.target[z] ])

                            if n == 0 and pno == 0:
                                sampweight = np.array([ lr 
                                                        for i, z in enumerate(ind) 
                                                        if (d-1)*subblocksize <= i and i < d*subblocksize
                                                        if logchoices[z] == train.target[z] ])
                            else:
                                sampweight = np.array([ lr * w * ocrl.n 
                                                           * max(0, qfunc(1, w, 1))
                                                        for i, z in enumerate(ind) 
                                                        if (d-1)*subblocksize <= i and i < d*subblocksize
                                                        if logchoices[z] == train.target[z]
                                                        for w in (1 / pchoices[z],)
                                                      ])

                            effn += sampweight.sum() / (lr * subblocksize)
                            effnsl += sampweight.sum() / (lr * subblocksize)

                            if np.any(y):
                                x = vstack([ v[i] 
                                             for i, z in enumerate(ind)
                                             if (d-1)*subblocksize <= i and i < d*subblocksize
                                             if logchoices[z] == train.target[z] ])
                                cls.partial_fit(x, y, classes=classes, sample_weight=sampweight)

                    print('{:<8d}\t{:<8.3f} [{:<10.3f}]\t{:<8.3f} [{:<10.3f}]\t{:<8.3f} [{:<10.3f}]\t{:<8.3f} [{:<10.3f}]'.format(
                                     loss.n, effn.mean(), effnsl.mean(), avtau.mean(), avtausl.mean(), loss.mean(), sincelast.mean(), logpv.mean(), logpvsl.mean()),
                               flush=True)   
                    sincelast = OnlineDRO.EasyAcc()
                    logpvsl = OnlineDRO.EasyAcc()
                    effnsl = OnlineDRO.EasyAcc()
                    avtausl = OnlineDRO.EasyAcc()

                    preds = cls.predict(testdocs)
                    ascores = []
                    for b in range(16):
                        bootie = np.random.RandomState(90210+b).choice(len(test.target), replace=True, size=len(test.target))
                        ascores.append(accuracy_score(y_true=test.target[bootie], y_pred=preds[bootie]))

                    print("test accuracy: {}".format(np.quantile(ascores, [0.05, 0.5, 0.95])))
            except KeyboardInterrupt:
                raise
            except:
                pass
    
OnlineDRO.flass()      

# + hidden=true


StarcoderdataJupyterScriptsDedupFiltered