#!/usr/bin/python
# Produces two polynomial matrices P1, P2
# P1 consists of first p1 columns of P
# P2 consists of remaining p2 columns of P

import pylab as pl
from polyfuncs import polsize

def colsplit(P, degP, p1, p2):
    P = pl.atleast_2d(P)
    if pl.atleast_2d(P).size is 0:
        P1 = pl.atleast_2d([])
        P2 = pl.atleast_2d([])
        return P1, None, P2, None

    rP, cP = polsize(P, degP)

    if not 0 <= p1 <= cP or not 0 <= p2 <= cP or p1+p2 != cP:
        print 'colsplit: Inconsistent numbers of columns'
        return None
        exit(1)
    
    P1 = pl.atleast_2d(pl.zeros((rP, (degP+1)*p1)))
    P2 = pl.atleast_2d(pl.zeros((rP, (degP+1)*p2)))
    
    
    for i in range(degP+1):
        P1[:,i*p1:(i+1)*p1] = P[:,i*cP:i*cP+p1]
        P2[:,i*p2:(i+1)*p2] = P[:,i*cP+p1:(i+1)*cP]
    
    return P1, degP, P2, degP

if __name__== "__main__":
    s = """P = pl.array([[1, 2, 3, 4, 1, 2, 3, 4],
           [9, 10, 11, 12, 9, 10, 11, 12],
           [5, 6, 7, 8, 5, 6, 7, 8]])"""
    print s
    exec(s)
    print """colsplit(P, 3, 1, 1)"""
    print colsplit(P, 3, 1, 1)