collatz_mod

import matplotlib.pyplot as plt
#from random import randrange

def itOne(n):
    return 1 if n==1 else (n*3+1)//2 if n%2 else n//2
    return 1 if n==1 else n*3+1 if n%2 else n//2

def itFull(n):
    ni=n
    cs=[ni]
    while ni>1:
        ni=itOne(ni)
        cs.append(ni)
    return cs

def pen(n):
    return n**1.0 if n>0 else 0
    ###alternatives
    ##return n**1.0 if n>0 else 0 #-(abs(n)**1.1)/1000

def gen(sp,c,r):
    if c==-10:
        aaa=1
    sc=[(float('inf'),0)]
    for i in range(1,len(r)):
        rg=sc[i-1][0]+pen(r[i])
        dg=sp[i-1][0]+pen(r[i]+c)
        sv = min( rg, dg )
        sd = int( rg > dg )
        ###alternatives
        ##sd = int( rg >= dg )        
        ##if rg==dg:
        ##    sd=randrange(2)
        sc+=[(sv,sd)]        
    return sc

def traceAndAdd(s,c,r):
    h=len(c)-1
    rn=r[:]
    for i in range(len(r)-1,-1,-1):
        if s[h][i][1] == 1:
            rn[i]+=c[h]
            h-=1
    return rn
                

def addFullIt(c,r):
    s=[ [(0,0)] ]
    for i in range(1,len(r)):
        s[0]+=[(s[0][-1][0]+pen(r[i]),0)]
    for i in range(0,len(c)-1):
        s+=[gen(s[i],c[i+1],r)]
    rn=traceAndAdd(s,c,r)
    return rn

n1=200#0
n2=400#0

a=list(range(n1,n2+1))
b=list(map(itFull,a))
c=sorted(b,key=lambda t:len(t),reverse=True)

r=[0]*((len(c[0])+1)*2)
cd=list(map(lambda l: list(map(lambda i,k:i-k, l[1:], l[:-1])), c))
rds=[]

for cde in cd:
    ###progress mark
    ##print('*',end='')
    cde=[0]+cde
    ro=r[:]
    r=addFullIt(cde,r)
    rd=list(map(lambda i,k:i-k, r,ro))
    rds+=[rd]

###reiteration, doesn't help much
##for cde,rde in zip(cd,rds):
##    print('*',end='')
##    cde=[0]+cde
##    r=list(map(lambda i,k:i-k, r,rde))
##    r=addFullIt(cde,r)


rf=[r[0]-sum(r)]
for i in range(1,len(r)):
    rf+=[rf[-1]+r[i]]

rps=sum(map(pen,r))
print('max: ',max(rf),', pen: ',rps)

plt.plot(rf)
plt.show()