
#/*##########################################################################
# Copyright (C) 2001-2013 European Synchrotron Radiation Facility
#
#              PyHST2  
#  European Synchrotron Radiation Facility, Grenoble,France
#
# PyHST2 is  developed at
# the ESRF by the Scientific Software  staff.
# Principal author for PyHST2: Alessandro Mirone.
#
# This program is free software; you can redistribute it and/or modify it 
# under the terms of the GNU General Public License as published by the Free
# Software Foundation; either version 2 of the License, or (at your option) 
# any later version.
#
# PyHST2 is distributed in the hope that it will be useful, but WITHOUT ANY
# WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
# FOR A PARTICULAR PURPOSE.  See the GNU General Public License for more
# details.
#
# You should have received a copy of the GNU General Public License along with
# PyHST2; if not, write to the Free Software Foundation, Inc., 59 Temple Place,
# Suite 330, Boston, MA 02111-1307, USA.
#
# PyHST2 follows the dual licensing model of Trolltech's Qt and Riverbank's PyQt
# and cannot be used as a free plugin for a non-free program. 
#
# Please contact the ESRF industrial unit (industry@esrf.fr) if this license 
# is a problem for you.
#############################################################################*/

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import subprocess as sub
import sys

import string
    ## in un futuro in cui potessero esserci piu di due cpu i cores potrebbero essere
    ## sur intervalli non contigui. In questo caso non bastera avere il numero totale di core
    ## e il primo numero. Ci vorranno piu range , ognuno individuato dall'inizio 
    ## e dal numero di cores  0,5,7,9-11

from . import string_six


    
def getCpuSetRange(cpuset_string, first=True):
    pos1 = cpuset_string.find( ":")
    pos2 = max(cpuset_string.rfind( "-"), cpuset_string.rfind( ","))
    if pos1 == -1 :
        raise Exception( "ERROR : wrong  output  in  cpuset_string " )
    if pos2 == -1 :
        if first:
            print( " **********************************!!!!!!!!!!!!!!!!!!!!!!!!!!\n"*10)
            print( " ERROR : you must request several OAR cores. PyHst cannot run on one core only , cpuset_string was  "+cpuset_string)
            raise Exception( "ERROR : you must request several OAR cores. PyHst cannot run on one core only , cpuset_string was "+cpuset_string )
        else:
            return []
    result = []
    msg = cpuset_string[pos1+1:]
    msgs = msg.split(",")
    for tok in msgs:
        ## print( tok)
        if "-" not in tok:
            result.append(  string.atoi(tok)  )
        else:
            a,b = map(string.atoi, tok.split("-"))
            result.extend( range(a,b+1))
    return result

def getCoresOrdered():
    
 
    comando = 'lscpu'
    if (sys.version_info >= (3, 7)):
        p1 = sub.Popen(args=comando.split( " ") ,stdin=sub.PIPE,stdout=sub.PIPE,stderr=sub.PIPE, text=True)
    elif (sys.version_info > (3, 0)):
        p1 = sub.Popen(args=comando.split( " ") ,stdin=sub.PIPE,stdout=sub.PIPE,stderr=sub.PIPE, universal_newlines=True)
    else:
        p1 = sub.Popen(args=comando.split( " ") ,stdin=sub.PIPE,stdout=sub.PIPE,stderr=sub.PIPE)
    # comando = 'grep node'
    # p2 = sub.Popen(args=comando.split( " ") ,stdin=p1.stdout,stdout=sub.PIPE,stderr=sub.PIPE)
    # comando = 'grep CPU'
    # p3 = sub.Popen(args=comando.split( " ") ,stdin=p2.stdout,stdout=sub.PIPE,stderr=sub.PIPE)
    # msg,errors=  p3.communicate()

    msg,errors=  p1.communicate()

    ## print( "  lscpu :  ",  msg)
    res = getCoresOrdered_fromstring(msg)
    return res

def     getCoresOrdered_fromstring(msg):
    msgs = msg.split("\n")
    count=0
    for l in msgs:
        count+=1
        # print( l)
        if "NUMA" in l:
            ncpus = string.atoi(l.split(":")[1])
            break
    else:
        raise Exception( "cannot understand output from lscpu ")



    result = []
    first= True
    for t in msgs[count:]:
        # if "CPU" in t:
        if "NUMA" in t:
            res = getCpuSetRange(t, first = first)
            first=False
            print ( " RES " , res)
            if(len(res)):
                result = result + res
                ncores4cpu = len(res)

    return ncores4cpu, result

if __name__=="__main__":
    print( getCpuSetRange(" asdasdas : 0,5,7,9-11"))
    s="""
NUMA node0 CPU(s):     0,2,4,6,8,10
NUMA node1 CPU(s):     1,3,5,7,9,11
"""
    print( getCoresOrdered_fromstring(s))
