# Install MOSEK if not already installed
!pip install mosek

# In this example we use the Fusion API
#    https://docs.mosek.com/latest/pythonfusion/index.html
# One could also use the lower level Optimizer API for Python
#    https://docs.mosek.com/latest/pythonapi/index.html
from mosek.fusion import *
import sys
import numpy as np

# Set up a very simple linear regression problem for demonstration
def mosek_least_squares(m, n, A, b, log):
  M = Model()
  x = M.variable(n)
  t = M.variable()

  # t >= ||Ax-b||_2
  M.constraint(Expr.vstack(t, Expr.sub(Expr.mul(A, x), b)),
               Domain.inQCone())
  
  M.objective(ObjectiveSense.Minimize, t)

  # Use the online demo optimization server
  M.optserverHost("http://solve.mosek.com:30080") 
  M.setLogHandler(sys.stdout if log else None)
  M.solve()

  print("Problem status: {prosta}".format(prosta=M.getProblemStatus()))
  print("residual: {t}".format(t=t.level()[0]))

  if M.getProblemStatus() == ProblemStatus.PrimalAndDualFeasible:
    return x.level()

# Random example
m, n = 20, 15
A = np.random.randn(m, n)
b = np.random.randn(m)

x = mosek_least_squares(m, n, A, b, False)

# 2-dimensional example with intercept
m = 30
X = np.random.rand(m)
Y = 3 * X + 1 + 0.5*(np.random.rand(m)-0.5)

a = mosek_least_squares(m, 2, np.vstack((X,np.ones(m))).transpose(), Y, False)

print("Regression line: {0:.3f} x + {1:.3f}".format(a[0], a[1]))
%matplotlib inline
import matplotlib.pyplot as plt
d = np.linspace(0, 1, 2)
plt.scatter(X, Y)
plt.plot(d, a[0] * d + a[1])